fix errs

SciML · Jan 4, 2024 · 34bff41 · 34bff41
1 parent bf77f69
commit 34bff41
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 6 deletions.
diff --git a/lib/OptimizationOptimisers/test/runtests.jl b/lib/OptimizationOptimisers/test/runtests.jl
@@ -61,7 +61,7 @@ using Zygote
 
         prob = OptimizationProblem(optprob, x0, _p)
         function callback(state, l)
-            Optimisers.adjust!(state.solver_state, 0.1/state.iter)
+            Optimisers.adjust!(state.solver_state, 0.1/state.iteration)
             return false
         end
         sol = solve(prob, Optimisers.Adam(0.1), maxiters = 1000, progress = false, callback = callback)

diff --git a/test/diffeqfluxtests.jl b/test/diffeqfluxtests.jl
@@ -34,14 +34,12 @@ function loss_adjoint(p)
 end
 
 iter = 0
-callback = function (p, l, pred)
-    global iter
-    iter += 1
+callback = function (state, l, pred)
 
     display(l)
 
     # using `remake` to re-create our `prob` with current parameters `p`
-    remade_solution = solve(remake(prob_ode, p = p), Tsit5(), saveat = tsteps)
+    remade_solution = solve(remake(prob_ode, p = state.u), Tsit5(), saveat = tsteps)
 
     # Tell sciml_train to not halt the optimization. If return true, then
     # optimization stops.

diff --git a/test/minibatch.jl b/test/minibatch.jl
@@ -18,7 +18,7 @@ function dudt_(u, p, t)
     ann(u, p, st)[1] .* u
 end
 
-callback = function (p, l, pred, args...; doplot = false) #callback function to observe training
+callback = function (state, l, pred, args...; doplot = false) #callback function to observe training
     display(l)
     # plot current prediction against data
     if doplot