diff --git a/Project.toml b/Project.toml
index 451430e..4be739d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,7 +31,7 @@ Distributions = "0.25.109"
 Flux = "0.12, 0.13, 0.14"
 LinearAlgebra = "1.7, 1.10"
 MLJBase = "1"
-MLJFlux = "0.5"
+MLJFlux = "0.5, 0.6"
 MLJModelInterface = "1.8.0"
 MLUtils = "0.4"
 Optimisers = "0.2, 0.3"