tpu_cores=8 not working

[rohitgr7]

🐛 Bug

After #2016 was fixed with PR #2033 the code is running perfectly on single tpu core and a specific tpu core but now not working with 8 tpu cores. After the training is complete getting RuntimeError: Cannot replicate if number of devices (1) is different from 8.

To Reproduce

Colab notebook

Expected behavior

Should train with 8 tpu cores with no error just like it works in case of a single core.

Environment

• pytorch/xla: nightly
• pytorch-lightning: master
• PyTorch Version (e.g., 1.0): 1.5
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): pip
• Python version: 3.7

[Borda]

To Reproduce
Colab notebook

I am sorry, I can't run it as it is read-only

@rohitgr7 mind check tests in #2094 as all the test there are passing, see: #2094 (comment)

[lezwon]

I'm seeing this issue when I run the notebook. Unsure of why it's happening. It's not consistent either but occurs the majority of times. I ran the code with an old release of lightning too, it still breaks. There is some change somewhere which is breaking it. If I run just the model with plain xmp.spawn , it seems to work fine.

[rohitgr7]

I have checked this on colab only. Will try it on kaggle-kernels and check if it's happening there too.

[williamFalcon]

@lezwon is this still an issue? Will reopen if it's still an issue since this is now tested

[lezwon]

I just ran the notebook. The training seems to be working fine, but .test() fails. Also calling .fit() the second time seems to throw the same error @rohitgr7 mentioned.