1 files changed, 3 insertions, 6 deletions

diff --git a/training/functional.py b/training/functional.py
index 3f5fa7e..e7c4320 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -375,7 +375,6 @@ def train_loop(
     num_val_steps = num_val_steps_per_epoch * num_epochs
 
     global_step = 0
-    train_step = 0
 
     avg_loss = AverageMeter()
     avg_acc = AverageMeter()
@@ -439,11 +438,11 @@ def train_loop(
                     loss, acc, bsz = loss_step(step, batch)
                     loss /= gradient_accumulation_steps
 
+                    accelerator.backward(loss)
+
                     avg_loss.update(loss.detach_(), bsz)
                     avg_acc.update(acc.detach_(), bsz)
 
-                    accelerator.backward(loss)
-
                     logs = {
                         "train/loss": avg_loss.avg.item(),
                         "train/acc": avg_acc.avg.item(),
@@ -455,9 +454,7 @@ def train_loop(
 
                     local_progress_bar.set_postfix(**logs)
 
-                    train_step += 1
-
-                    if train_step % gradient_accumulation_steps == 0:
+                    if ((step + 1) % gradient_accumulation_steps == 0) or ((step + 1) == len(train_dataloader)):
                         on_before_optimize(lr_scheduler.get_last_lr()[0], epoch)
 
                         optimizer.step()