From e7d77b38e95584bc4a2b595197d4e0de61632604 Mon Sep 17 00:00:00 2001
From: "L. Srinivasa Mohan" <lsmohan@gmail.com>
Date: Tue, 7 Jan 2025 16:16:41 +0530
Subject: [PATCH] Update OptimizationOptimisers.jl

Modify check for last iteration to set the parameters to the best minimizer if save_best if enabled.
---
 lib/OptimizationOptimisers/src/OptimizationOptimisers.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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,