import ParameterSchedulers to avoid name conflict with Scheduler

FluxML · Mar 9, 2024 · 3ac6a9c · 3ac6a9c
1 parent 0c2fdf5
commit 3ac6a9c
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/src/FluxTraining.jl b/src/FluxTraining.jl
@@ -22,7 +22,8 @@ using Statistics: mean
 using UUIDs
 using Zygote
 using ChainRulesCore
-using ParameterSchedulers
+import ParameterSchedulers
+import ParameterSchedulers: Sequence, Shifted, Sin
 using TensorBoardLogger: TBLogger, log_value, log_image, log_text, log_histogram, tb_overwrite
 using Zygote: Grads, gradient
 using ValueHistories

diff --git a/test/callbacks/scheduler.jl b/test/callbacks/scheduler.jl
@@ -27,7 +27,7 @@ include("../imports.jl")
     end
 
     @testset "Regression test for #122" begin
-        learner = testlearner(Recorder(), Scheduler(LearningRate => ParameterSchedulers.CosAnneal(λ0=0,λ1=0,period=0)), ToGPU())
+        learner = testlearner(Recorder(), Scheduler(LearningRate => CosAnneal(λ0=0,λ1=0,period=0)), ToGPU())
         @test_nowarn FluxTraining.Graphs.topological_sort_by_dfs(learner.callbacks.graph)
     end
 end
diff --git a/test/imports.jl b/test/imports.jl
@@ -1,7 +1,7 @@
 using ReTest
 import Optimisers
 using FluxTraining
-using ParameterSchedulers
+import ParameterSchedulers: Interpolator, Poly, CosAnneal
 using Colors
 using ImageIO
 using FluxTraining: protect, Events, Phases, runstep, runepoch, epoch!, step!, testlearner, SanityCheckException