Longformer encdec fails on TPU with "scalar type not supported" error

Hi,

I'm trying to train longformer encdec on Cloud TPUs with the following settings.
- GCE instance:
`e2-standard-32 (128GB RAM)`
- GCE image family and OS:
`ubuntu-2004-lts (Ubuntu20.04LTS)`
- Docker container on GCE instance:
`gcr.io/tpu-pytorch/xla:r1.6 (Latest stable container distributed at Pytorch/XLA repo)`
- TPU instance:
`v3-8 (with --version=pytorch-1.6 option)`
- Lib versions on the docker container
- python: `3.6.10`
- pytorch: `1.6.0a0+b31f58d`
- pytorch XLA: `1.6`
- pytorch-lightning: `0.8.5`

and script/summarization.py options:
```python scripts/summarization.py
--model_path="./longformer-encdec-8192"
--max_input_len=8192
--grad_ckpt
--batch_size=1
--gpus=0
```

Setting `tpu_cores=8` to pl.trainer loads the converted model on TPU with no error, but then fails after sanity check starts:
```
Namespace(adafactor=False, attention_dropout=0.1, attention_mode='sliding_chunks', attention_window=512, batch_size=1, debug=False, disable_checkpointing=False, epochs=5, fp32=False, gpus=0, grad_accum=1, grad_ckpt=True, label_smoothing=0.0, lr=3e-05, max_input_len=16384, max_output_len=256, model_path='./longformer-encdec-16384', no_progress_bar=False, num_workers=0, resume_ckpt=None, save_dir='summarization', save_prefix='test', seed=1234, test=False, tokenizer='facebook/bart-base', val_every=1.0, val_percent_check=1.0, warmup=1000)
GPU available: False, used: False
TPU available: True, using: 8 TPU cores
Using native 16bit precision.
training on 8 TPU cores
INIT TPU local core: 0, global rank: 0
INIT TPU local core: 1, global rank: 1
INIT TPU local core: 2, global rank: 2
INIT TPU local core: 3, global rank: 3
INIT TPU local core: 4, global rank: 4
INIT TPU local core: 5, global rank: 5
INIT TPU local core: 6, global rank: 6
INIT TPU local core: 7, global rank: 7

| Name | Type | Params
---------------------------------------------------------------------------
0 | model | LongformerEncoderDecoderForConditionalGeneration | 459 M
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your val_dataloader has `shuffle=True`, it is best practice to turn this off for validation and test dataloaders.
warnings.warn(*args, **kwargs)
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 32 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
warnings.warn(*args, **kwargs)
Validation sanity check: 0it [00:00, ?it/s]2020-10-05 02:34:02.456683: E 657 tensorflow/compiler/xla/xla_client/tf_logging.cc:11] Check failed: scalar_value.isIntegral()

```
**Note: I will paste full stack-trace just after this post**

I have no idea to debug this error, but it seems to come from XLA multiprocessing or using docker...?

Any comments and suggestions is appreciated.