Problem trying to run AlbertForMaskedLM on Colab TPU: TypeError: can't pickle torch._C.ScriptFunction objects when calling xm.send_cpu_data_to_device(model, dev) 
Author: ayanyuegupta
Pull request: no

# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): AlbertForMaskedLM

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [X ] the official example scripts: (give details below)
* [X ] 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

Steps to reproduce the behavior:

Related issues: https://github.com/pytorch/xla/issues/1909

The following has already been talked about here (https://github.com/huggingface/transformers/pull/3743) but I couldn't find a solution? Apologies if I'm posting about something that's already been dealt with, pretty new to all of this.

I am running the following code on colab on a TPU session taken from the example here: https://huggingface.co/transformers/model_doc/albert.html#albertformaskedlm

```
import os
import torch
import torch_xla
import torch_xla.core.xla_model as xm

assert os.environ['COLAB_TPU_ADDR']

dev = xm.xla_device()

from transformers import AlbertTokenizer, AlbertForMaskedLM
import torch

tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
model = AlbertForMaskedLM.from_pretrained('albert-base-v2')
model = xm.send_cpu_data_to_device(model, dev)
model = model.to(dev)
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1

data = input_ids.to(dev)

outputs = model(data, masked_lm_labels=data)
loss, prediction_scores = outputs[:2]
```

I haven't done anything to the example code except move ```input_ids``` and ```model``` onto the TPU device using ```.to(dev)``` and ```xm.send_cpu_data_to_device(model, dev)```. It seems everything is moved to the TPU no problem as when I input ```data``` I get the following output: ```tensor([[    2, 10975,    15,    51,  1952,    25, 10901,     3]], device='xla:1')```

However when I run this code I get the following error:

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-6-b7b68efc9620> in <module>()
     11 tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
     12 model = AlbertForMaskedLM.from_pretrained('albert-base-v2')
---> 13 model = xm.send_cpu_data_to_device(model, dev)
     14 model = model.to(dev)
     15 input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1

18 frames
/usr/lib/python3.6/copy.py in copy(x)
     94         reductor = getattr(x, "__reduce_ex__", None)
     95         if reductor:
---> 96             rv = reductor(4)
     97         else:
     98             reductor = getattr(x, "__reduce__", None)

TypeError: can't pickle torch._C.ScriptFunction objects
```
Anyone know what's going on?

## Expected behavior

I expected the AlbertForMaskedLM model to work on colab TPU without any errors.

## 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.8.0
- Platform: Linux-4.19.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.5.0a0+ab660ae (False)
- Tensorflow version (GPU?): 2.2.0-rc3 (False)
- Using GPU in script?: no, attempting to use TPU
- Using distributed or parallel set-up in script?: no


Comments:
Seems fixed now--delete transformers installed via pip and install by cloning this repo.