undefined symbol when importing PyTorch Lightning after installing XLA
🐛 Bug

Following the doc, in order to use TPU on colab pro, I can't import PyTorch Lightning after installing XLA:
ImportError: /usr/local/lib/python3.6/dist-packages/_XLAC.cpython-36m-x86_64-linux-gnu.so: undefined symbol: _ZN2at13nanmedian_outERNS_6TensorES1_RKS0_lb

To Reproduce
https://colab.research.google.com/drive/1dv21HCCIMdv2r9iSqSRpeOmOJ7Hh8TXF?usp=sharing
Expected behavior
Well, I should be able to import pytorch_lightning I guess :)
It worked few days ago.
Installing pytorch_lightning before XLA leads to an error telling that XLA isn't installed when calling trainer.fit.

Environment

PyTorch Version (e.g., 1.0): 1.7.0 & 1.8.0
OS (e.g., Linux): Linux (colab)
How you installed PyTorch (conda, pip, source): Already installed (colab)
Build command you used (if compiling from source): N/A
Python version: 3.6
CUDA/cuDNN version: 10.1
GPU models and configuration: TPU
Any other relevant information:

Additional context