Error in run_language_modeling on TPU: Transferring data with element type U8 has not been implemented on TPUs
## 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: 3.1.0
- Platform: Linux-4.19.0-10-cloud-amd64-x86_64-with-debian-9.13
- Python version: 3.7.7
- PyTorch version (GPU?): 1.7.0a0+626e410 (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: No (TPU)
- Using distributed or parallel set-up in script?: No

### Who can help
 albert, bert, GPT2, XLM: @LysandreJik 

<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
 albert, bert, GPT2, XLM: @LysandreJik 
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Speed and Memory Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 Translation: @sshleifer
 Summarization: @sshleifer
 TextGeneration: @TevenLeScao 
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @TevenLeScao
 blenderbot: @mariamabarham
 Bart: @sshleifer
 Marian: @sshleifer
 T5: @patrickvonplaten
 Longformer/Reformer: @patrickvonplaten
 TransfoXL/XLNet: @TevenLeScao 
 examples/seq2seq: @sshleifer
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 -->

## Information

Model I am using (Bert, XLNet ...): GPT2 - (Not pretrained) 

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

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

## To reproduce
I am trying to train a GPT model from scratch with TPU v2-8 in GCP. I have tried several setups but consistently getting the error below. I have simplified the arguments as much as I can do to increase the reproducibility.

Steps to reproduce the behavior:

1. Pytorch/XLA installed through docker: gcr.io/tpu-pytorch/xla:nightly_3.7
2. Transformers 3.1.0 installed through pip.
3. I used the latest xla_spawn.py with previous version of [run_language_modeling.py](https://github.com/huggingface/transformers/blob/a75c64d80c76c3dc71f735d9197a4a601847e0cd/examples/language-modeling/run_language_modeling.py) because the latest version of run_language_modeling gives error about cache_dir.
4. Simple txt file with 10k lines given as input.


```bash
python xla_spawn.py --num_cores=1 run_language_modeling.py --output_dir=/home/GPT/output --model_type=gpt2 --tokenizer_name=gpt2 --do_train --train_data_file=/home/GPT/corpus/oscar_wiki_opus_eval.txt --per_device_train_batch_size=128
```

Output:

```bash
09/08/2020 11:28:50 - WARNING - run_language_modeling -   Process rank: -1, device: xla:1, n_gpu: 0, distributed training: False, 16-bits training: False
09/08/2020 11:28:50 - INFO - run_language_modeling -   Training/evaluation parameters TrainingArguments(output_dir='/home/GPT/output', overwrite_output_dir=False, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, prediction_loss_only=False, per_device_train_batch_size=128, per_device_eval_batch_size=8, per_gpu_train_batch_size=None, per_gpu_eval_batch_size=None, gradient_accumulation_steps=1, learning_rate=5e-05, weight_decay=0.0, adam_beta1=0.9, adam_beta2=0.999, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=3.0, max_steps=-1, warmup_steps=0, logging_dir='runs/Sep08_11-28-39_b0b61c9b1d63', logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=1, tpu_metrics_debug=False, debug=False, dataloader_drop_last=False, eval_steps=1000, past_index=-1, run_name=None, disable_tqdm=False, remove_unused_columns=True)
09/08/2020 11:28:50 - WARNING - run_language_modeling -   You are instantiating a new config instance from scratch.
09/08/2020 11:28:50 - INFO - run_language_modeling -   Training new model from scratch
/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/transformers/modeling_auto.py:732: FutureWarning: The class `AutoModelWithLMHead` is deprecated and will be removed in a future version. Please use `AutoModelForCausalLM` for causal language models, `AutoModelForMaskedLM` for masked language models and `AutoModelForSeq2SeqLM` for encoder-decoder models.
  FutureWarning,
/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/transformers/tokenization_utils_base.py:1321: FutureWarning: The `max_len` attribute has been deprecated and will be removed in a future version, use `model_max_length` instead.
  FutureWarning,
09/08/2020 11:28:54 - INFO - filelock -   Lock 140191457088144 acquired on /home/GPT/corpus/cached_lm_GPT2Tokenizer_1024_oscar_wiki_opus_eval.txt.lock
09/08/2020 11:29:09 - INFO - filelock -   Lock 140191457088144 released on /home/GPT/corpus/cached_lm_GPT2Tokenizer_1024_oscar_wiki_opus_eval.txt.lock
Traceback (most recent call last):
  File "xla_spawn.py", line 72, in <module>
    main()
  File "xla_spawn.py", line 68, in main
    xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 387, in spawn
    _start_fn(0, pf_cfg, fn, args)
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 324, in _start_fn
    fn(gindex, *args)
  File "/home/GPT/run_language_modeling.py", line 294, in _mp_fn
    main()
  File "/home/GPT/run_language_modeling.py", line 252, in main
    prediction_loss_only=True,
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/transformers/trainer.py", line 229, in __init__
    self.model = model.to(args.device) if model is not None else None
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 612, in to
    return self._apply(convert)
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  [Previous line repeated 1 more time]
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 402, in _apply
    self._buffers[key] = fn(buf)
  File "/root/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 610, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:383 : Check failed: session->session()->Run( session_work->feed_inputs, session_work->outputs_handles, &outputs) == ::tensorflow::Status::OK() (Unimplemented: From /job:tpu_worker/replica:0/task:0:
Attempted to transfer array of shape u8[1,1,1024,1024] to a TPU device. Transferring data with element type U8 has not been implemented on TPUs.
	 [[{{node XRTAllocateFromTensor_6}}]] vs. OK)
*** Begin stack trace ***
	tensorflow::CurrentStackTrace()

	xla::util::MultiWait::Complete(std::function<void ()> const&)



	clone
*** End stack trace ***

terminate called after throwing an instance of 'std::runtime_error'
  what():  tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:1110 : Check failed: session->session()->Run( feed_inputs, {}, {cached_node.operations[0]}, &outputs) == ::tensorflow::Status::OK() (Not found: Op type not registered 'XRTMemoryInfo' in binary running on n-cec96fe8-w-0. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed. vs. OK)
*** Begin stack trace ***
	tensorflow::CurrentStackTrace()
	xla::XrtComputationClient::ReleaseHandles(std::vector<xla::XrtComputationClient::DeviceHandle, std::allocator<xla::XrtComputationClient::DeviceHandle> >*, std::function<xla::XrtSession::CachedNode const& (xla::XrtSession*, tensorflow::Scope const&, std::string const&)> const&, xla::metrics::Metric*, xla::metrics::Counter*)
	xla::XrtComputationClient::HandleReleaser()
	xla::util::TriggeredTask::Runner()


	clone
*** End stack trace ***

Aborted (core dumped)
```

## Expected behavior
I would expect it to train GPT model on TPU with given txt file from scratch.

P.S.:
I have also tried with the stable version of XLA (gcr.io/tpu-pytorch/xla:r1.6) and different model_type and tokenizer_name (bert). I am getting the same error in both cases. 

I have tried several things to change the input type to float in the source code of transformers, but I was not successful. Any tips for that direction would be good for me, too. 
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
