tests pass now pls

SciML · Sep 13, 2024 · 2a803ff · 2a803ff
1 parent 6e1999d
commit 2a803ff
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/Project.toml b/Project.toml
@@ -30,7 +30,7 @@ LinearAlgebra = "1.10"
 Logging = "1.10"
 LoggingExtras = "0.4, 1"
 MLUtils = "0.4.4"
-OptimizationBase = "2.0.2"
+OptimizationBase = "2.0.3"
 Printf = "1.10"
 ProgressLogging = "0.1"
 Reexport = "1.2"

diff --git a/src/sophia.jl b/src/sophia.jl
@@ -1,4 +1,4 @@
-using Optimization.LinearAlgebra
+using Optimization.LinearAlgebra, MLUtils
 
 struct Sophia
     η::Float64
@@ -80,14 +80,14 @@ function SciMLBase.__solve(cache::OptimizationCache{
     for _ in 1:maxiters
         for (i, d) in enumerate(data)
             if cache.f.fg !== nothing && dataiterate
-                x = cache.f.fg(G, θ, d)
+                x = cache.f.fg(gₜ, θ, d)
             elseif dataiterate
-                cache.f.grad(G, θ, d)
+                cache.f.grad(gₜ, θ, d)
                 x = cache.f(θ, d)
             elseif cache.f.fg !== nothing
-                x = cache.f.fg(G, θ)
+                x = cache.f.fg(gₜ, θ)
             else
-                cache.f.grad(G, θ)
+                cache.f.grad(gₜ, θ)
                 x = cache.f(θ)
             end
             opt_state = Optimization.OptimizationState(; iter = i,

diff --git a/test/minibatch.jl b/test/minibatch.jl
@@ -19,7 +19,7 @@ function dudt_(u, p, t)
     ann(u, p, st)[1] .* u
 end
 
-function callback(state, l, pred) #callback function to observe training
+function callback(state, l) #callback function to observe training
     display(l)
     return false
 end