🐛 run_ner.py runtime error linked to TPU training
Author: vinmorel
Pull request: no

# 🐛 Bug

## Information

Model I am using **Longformer For Token Classification**

Language I am using the model on **German**:

The problem arises when using:
* [x] the official example scripts:
T**he problem arises when trying to run run_ner.py on google colab in TPU fp16 mode.**  


The tasks I am working on is:
* [x] an official GLUE/SQUaD task:
**CoNLL NER**


## To reproduce


I have a colab up if you want to see exactly what I did. You just need to upload data files in a new folder in /Content/<YOUR_FOLDER> to use in training and make sure to modify the run_ner paths correspondingly. [Google Colab](https://github.com/vinmorel/transformers/blob/master/run_ner_TPU.ipynb)


<!-- 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.-->

Runtime error after running the following :
```
!python3 "/content/transformers/examples/token-classification/run_ner.py" --data_dir "/content/CoNLL/" \
--labels "/content/CoNLL/labels.txt" \
--model_name_or_path "allenai/longformer-base-4096" \
--output_dir "xlnet-base-cased" \
--max_seq_length  200 \
--num_train_epochs 2 \
--per_device_train_batch_size 1 \
--save_steps 750 \
--seed 1 \
--do_train \
--do_eval  \
--do_predict \
--fp16
```
```
Traceback (most recent call last):
  File "/content/transformers/examples/token-classification/run_ner.py", line 303, in <module>
    main()
  File "/content/transformers/examples/token-classification/run_ner.py", line 228, in main
    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer.py", line 390, in train
    model, optimizer = amp.initialize(model, optimizer, opt_level=self.args.fp16_opt_level)
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/frontend.py", line 358, in initialize
    return _initialize(models, optimizers, _amp_state.opt_properties, num_losses, cast_model_outputs)
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/_initialize.py", line 171, in _initialize
    check_params_fp32(models)
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/_initialize.py", line 93, in check_params_fp32
    name, param.type()))
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/_amp_state.py", line 32, in warn_or_err
    raise RuntimeError(msg)
RuntimeError: Found param longformer.embeddings.word_embeddings.weight with type torch.FloatTensor, expected torch.cuda.FloatTensor.
When using amp.initialize, you need to provide a model with parameters
located on a CUDA device before passing it no matter what optimization level
you chose. Use model.to('cuda') to use the default device.
```

## Expected behavior

Would expect the code to run and start training on TPU without the runtime error.

## 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.11.0
- Platform: Linux-4.19.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.6.0.dev20200528 (False)
- Tensorflow version (GPU?): 2.2.0 (False)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No


Comments:
Hi! did you solve your issue?