Do not prepare lr scheduler as it as the right number of steps (#24088)

* Do not prepare lr scheduler as it as the right number of steps * Trigger CI * Trigger CI * Trigger CI * Add fake comment * Remove fake comment * Trigger CI please!
huggingface · Jun 7, 2023 · 2c887cf · 2c887cf
1 parent 12298cb
commit 2c887cf
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -1747,9 +1747,7 @@ def _inner_training_loop(
 
         # prepare using `accelerator` prepare
         if use_accelerator_prepare:
-            model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
-                self.model, self.optimizer, self.lr_scheduler
-            )
+            model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
 
         if self.is_fsdp_enabled:
             self.model = model
@@ -1996,6 +1994,7 @@ def _inner_training_loop(
                         optimizer_was_run = scale_before <= scale_after
                     else:
                         self.optimizer.step()
+                        optimizer_was_run = not self.accelerator.optimizer_step_was_skipped
 
                     if optimizer_was_run:
                         # Delay optimizer scheduling until metrics are generated