

TITLE
Generate text with `model.generate` on TPU does not work

COMMENTS
4

REACTIONS
+1: 1
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## Environment info
- `transformers` version: 4.7.0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29 (Ubuntu 20.04.2 LTS)
- Python version: 3.8.5
- PyTorch version (GPU?): 1.8.1+cu102 (False)
- PyTorch XLA version: 1.8.1
- Using GPU in script?: No, using TPU
- Using distributed or parallel set-up in script?: No, using a single TPU core

### Who can help
@patrickvonplaten

## Information

Model I am using (Bert, XLNet ...): `facebook/m2m100_1.2B`, but other text generating models have the same problem.

The problem arises when using:
* [ ] 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

On a machine with a TPU run:

```python
import torch_xla.core.xla_model as xm
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

model_name = 'facebook/m2m100_1.2B'
source_lang = 'en'
target_lang = 'de'

docs = [
    "This is some document to translate.",
    "And another document to translate."
]

device = xm.xla_device()

model = M2M100ForConditionalGeneration.from_pretrained(model_name).to(device)

tokenizer = M2M100Tokenizer.from_pretrained(model_name, src_lang=source_lang)
encoded_docs = tokenizer(docs, return_tensors='pt', padding=True).to(device)

generated_tokens = model.generate(**encoded_docs, forced_bos_token_id=tokenizer.get_lang_id(target_lang))
```

The call to `model.generate()` runs without ever terminating. It seems to be stuck somewhere in the beam search.

The same code runs perfectly fine on CPUs and GPUs.

## Expected behavior

I'd expect that the generation of text works in the same way as for CPUs and GPUs.
