From b0f2a3f8fa786a7b189bcd8b86cd0f3c0cc89b3a Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 6 Oct 2024 22:42:25 -0400 Subject: [PATCH] fix: match iterations count in Sophia --- src/sophia.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/sophia.jl b/src/sophia.jl index 88b0812c3..b63f0c099 100644 --- a/src/sophia.jl +++ b/src/sophia.jl @@ -75,7 +75,7 @@ function SciMLBase.__solve(cache::OptimizationCache{ gₜ = zero(θ) mₜ = zero(θ) hₜ = zero(θ) - for _ in 1:maxiters + for epoch in 1:maxiters for (i, d) in enumerate(data) if cache.f.fg !== nothing && dataiterate x = cache.f.fg(gₜ, θ, d) @@ -88,7 +88,7 @@ function SciMLBase.__solve(cache::OptimizationCache{ cache.f.grad(gₜ, θ) x = cache.f(θ) end - opt_state = Optimization.OptimizationState(; iter = i, + opt_state = Optimization.OptimizationState(; iter = i + (epoch - 1) * length(data), u = θ, objective = first(x), grad = gₜ,