fix iterator overflow when gradient accumulation is 1 (#35960)

huggingface · Jan 29, 2025 · 7547f55 · 7547f55
1 parent 4d3b107
commit 7547f55
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -2434,6 +2434,8 @@ def _inner_training_loop(
                 remainder = args.gradient_accumulation_steps
             update_step = -1
             total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1
+            if args.gradient_accumulation_steps == 1:
+                total_updates -= 1
             for _ in range(total_updates):
                 update_step += 1
                 num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder