diff --git a/Project.toml b/Project.toml
index e7b1ce9de..c41e46153 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "FluxTraining"
 uuid = "7bf95e4d-ca32-48da-9824-f0dc5310474f"
 authors = ["lorenzoh <lorenz.ohly@gmail.com>"]
-version = "0.3.9"
+version = "0.3.10"
 
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
@@ -40,7 +40,7 @@ ImageCore = "0.8, 0.9, 0.10"
 InlineTest = "0.2"
 OnlineStats = "1.5"
 Optimisers = "0.2, 0.3"
-ParameterSchedulers = "0.3.1"
+ParameterSchedulers = "0.3.1, 0.4"
 Parameters = "0.12"
 PrecompileTools = "1"
 PrettyTables = "1, 1.1, 1.2, 2"
diff --git a/src/FluxTraining.jl b/src/FluxTraining.jl
index 436f32086..8b5aab1ea 100644
--- 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
index 9ab6b1e43..6b223caab 100644
--- 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
index 2ba5e4dcb..5b9da4a72 100644
--- 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