diff --git a/lib/OptimizationOptimisers/src/OptimizationOptimisers.jl b/lib/OptimizationOptimisers/src/OptimizationOptimisers.jl
index da2299e01..59a1d3520 100644
--- a/lib/OptimizationOptimisers/src/OptimizationOptimisers.jl
+++ b/lib/OptimizationOptimisers/src/OptimizationOptimisers.jl
@@ -130,12 +130,12 @@ function SciMLBase.__solve(cache::OptimizationCache{
                         min_err = x
                         min_θ = copy(θ)
                     end
-                    if i == length(data)*epochs  #Last iter, revert to best.
+                    if iterations == length(data)*epochs  #Last iter, revert to best.
                         opt = min_opt
                         x = min_err
                         θ = min_θ
                         cache.f.grad(G, θ, d)
-                        opt_state = Optimization.OptimizationState(iter = i,
+                        opt_state = Optimization.OptimizationState(iter = iterations,
                             u = θ,
                             objective = x[1],
                             grad = G,