[Trainer] Fix optimizer step on PyTorch TPU (#24389)

* update optimizer step for tpu * add comment

[Trainer] Fix optimizer step on PyTorch TPU (#24389)
* update optimizer step for tpu * add comment
1815d186 · Meghan Cowan · GitHub · 4c6e4295 · 1815d186
Unverified Commit 1815d186 authored Jun 21, 2023 by Meghan Cowan Committed by GitHub Jun 21, 2023
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

src/transformers/trainer.py src/transformers/trainer.py +2 -1

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1859,7 +1859,8 @@ class Trainer:
                            self.scaler.step(self.optimizer)
                            self.scaler.update()
                        else:
-                            xm.optimizer_step(self.optimizer)
+                            # tpu-comment: accelerate wrapped optimizers call xm.optimizer_step
+                            self.optimizer.step()
                    elif self.do_grad_scaling:
                        scale_before = self.scaler.get_scale()
                        self.scaler.step(self.optimizer)