Unable to train with 8 cores TPU
🐛 Bug
Unable to train with TPU cores = 8. It trained one full fold, and it failed during the next fold.
I have a validation loop and it still fails. I am trying on kaggle.
To Reproduce
Fork edit this notebook and run on TPU
https://www.kaggle.com/krisho007/fork-of-melanoma-with-pylightning
Additional context

CUDA:

GPU:
available:         False
version:           10.2


Packages:

numpy:             1.18.5
pyTorch_debug:     False
pyTorch_version:   1.5.0
pytorch-lightning: 0.8.5
tensorboard:       2.2.2
tqdm:              4.45.0


System:

OS:                Linux
architecture:

64bit



processor:         x86_64
python:            3.7.6
version:           #1 SMP Thu Aug 6 20:22:07 PDT 2020