Trainer and Colab TPU: Training loss isn't declining
Author: gustavscholin
Pull request: no

# 🐛 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


Comments:
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!