Training loss doesn't decrease on TPU while works fine on GPU
For the summarization task, where I train with an encoder-decoder model on GPU, it works fine and the loss gets lower over iterations. But when I change the device to `device = xm.xla_device()` and optimizer to `xm.optimizer_step(optimizer, barrier=True)` on single-core TPU, the training loss remains nearly unchanged!!

**Here's the reproducible code:** https://colab.research.google.com/drive/1pC2CF3ipWt0eJrfXdznwAZD3zs0sX1kd?usp=sharing
Is it a bug or I am missing something?
Maybe @sgugger has an idea or knows someone who does!<|||||>This is more of a question for the PyTorch XLA folks, since you're not using any of our tools for training.<|||||>> This is more of a question for the PyTorch XLA folks, since you're not using any of our tools for training.

Thanks I have asked them and here's the solution: https://github.com/pytorch/xla/issues/3675#issuecomment-1171702988<|||||>Ah yes, I missed it but it's indeed a common mistake on TPUs!
(FYI: by using Accelerate to power your training loop, your mistake would have been automatically fixed ;-) )<|||||>Nice to know! Thanks for the information @sgugger. Actually, I started with the HF Trainer but facing [this](https://github.com/huggingface/transformers/issues/14989#issue-1091070983) issue I moved to [this](https://github.com/huggingface/transformers/issues/14989#issuecomment-1003349939) solution which used PyTorch loop instead though I am using v4.18.0.