tests pass now pls

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

Hide file tree

Showing 2 changed files with 6 additions and 6 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,