TPU training doesn't work on Colab. No TPU devices were found

[sirgarfieldc]

After installing these two dependency following the doc

pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.8.1-cp37-cp37m-linux_x86_64.whl
conda install pytorch-lightning -c conda-forge

I get 2 erros

Traceback (most recent call last):
  File "/root/anaconda3/envs/my-rdkit-env/lib/python3.7/site-packages/pytorch_lightning/utilities/xla_device.py", line 47, in wrapper
    return queue.get_nowait()
  File "/root/anaconda3/envs/my-rdkit-env/lib/python3.7/multiprocessing/queues.py", line 126, in get_nowait
    return self.get(False)
  File "/root/anaconda3/envs/my-rdkit-env/lib/python3.7/multiprocessing/queues.py", line 107, in get
    raise Empty
_queue.Empty

This one above is encountered when importing pytorch_lightning. This error does not prevent importing the package

pytorch_lightning.utilities.exceptions.MisconfigurationException: No TPU devices were found

This is encountered when running training script

Google around a bit and someone mentioned restarting Colab runtime solve the isseu. I restart and the issue goes away for a bit. Not sure what difference restarting made.

Following a force quit of my training script.

2 errors mentioned above come back.

This time restart does not solve the problem. Now I can no longer run GPU training

Software versions

Python 3.7.10
torch 1.8.1
torch_xla 1.8.1
pytorch_lightning 1.2.10