diff --git a/docs/Project.toml b/docs/Project.toml
index c91c30efd..568495207 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -58,9 +58,9 @@ MLUtils = "0.4"
 NNlib = "0.9"
 Optimisers = "0.2, 0.3"
 Optimization = "3.9"
-OptimizationOptimJL = "0.1"
-OptimizationOptimisers = "0.1"
-OptimizationPolyalgorithms = "0.1"
+OptimizationOptimJL = "0.2"
+OptimizationOptimisers = "0.2"
+OptimizationPolyalgorithms = "0.2"
 OrdinaryDiffEq = "6.31"
 Plots = "1.36"
 Printf = "<0.0.1, 1"
diff --git a/docs/src/examples/neural_ode_weather_forecast.md b/docs/src/examples/neural_ode_weather_forecast.md
index 6af55fa4c..0ca50ba75 100644
--- a/docs/src/examples/neural_ode_weather_forecast.md
+++ b/docs/src/examples/neural_ode_weather_forecast.md
@@ -136,7 +136,7 @@ end
 
 function train(t, y, obs_grid, maxiters, lr, rng, p = nothing, state = nothing; kwargs...)
     log_results(ps, losses) = (p, loss) -> begin
-        push!(ps, copy(p))
+        push!(ps, copy(p.u))
         push!(losses, loss)
         false
     end