Trainer and Colab TPU: Training loss isn't declining
# 🐛 Bug
When using the Trainer with a Colab TPU it seems like the training loss remains more or less constant around 1.11 during the whole training process. With Colab GPU/CPU however it you get the expected loss decline.
## Information

Model I am using (Bert, XLNet ...):Bert Base Cased

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

When you run `run_glue.py` with `xla_spawn.py` as shown in Running on TPUs section in the examples README.

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

The problem is reproduced in this Colab notebook:
https://colab.research.google.com/drive/1PLPReml5UOHC3iQdx8QvfCxuhDekTD2l?usp=sharing

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

The training loss declines the same way it does when using GPU/CPU.

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 2.9.1
- Platform: Linux-4.19.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.6.0a0+96885f7 (False)
- Tensorflow version (GPU?): 2.2.0 (False)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

I meeting the same issue with a custom training script (built using GLUE script as model), using ELECTRA for classification task.<|||||>Hi! This should have been solved by https://github.com/huggingface/transformers/pull/4450.
Could you install from source and let me know if it fixes your issue?<|||||>The code is running, but I receive this warning :

```
UserWarning: This overload of addcdiv_ is deprecated:
	addcdiv_(Number value, Tensor tensor1, Tensor tensor2)
Consider using one of the following signatures instead:
	addcdiv_(Tensor tensor1, Tensor tensor2, *, Number value) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:760.)
  p.data.addcdiv_(-step_size, exp_avg, denom)
```

and it's **very slow** (2h for 1 epoch). 

On a K80 GPU, it takes 30 minutes...

---

Running on 8 cores is not working but I think this is due to the RAM limitations.<|||||>Can confirm that the training loss is now declining as expected in colab notebook and that one epoch takes around 2h per epoch with 1 TPU-core. I also tested with a K80 but unlike for @Colanim it took around 5h per epoch, which I find reasonable in the sense that the TPU is faster. The batch size was the default Trainer value for both tests (which is 5 I believe).

I get @Colanim's UserWarning regardless of using GPU or TPU.

So imo the bug is fixed!