tests pass now pls

Author: Vaibhavdixit02

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"

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,

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