diff --git a/Project.toml b/Project.toml
index d60fa60..150dd8d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -150,7 +150,7 @@ MCMCChains = "5.0, 6"
 MCMCDiagnosticTools = "0.1"
 MLJ = "0.19"
 MLJBase = "0.20, 0.21"
-MLJFlux = "0.2"
+MLJFlux = "0.2, 0.6"
 MLJXGBoostInterface = "0.2, 0.3"
 MacroTools = "0.5"
 Makie = "0.19"