Pytorch-lightning: undefined symbol when importing PyTorch Lightning after installing XLA

[clementpoiret]

🐛 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

[github-actions[bot]]

Hi! thanks for your contribution!, great first issue!

[clementpoiret]

Seems just a nightly issue, working well with torch XLA 1.7... Is there a specific purpose of using nightly builds as stated in PL's doc?

[Borda]

we want to get nightly back, just have not done yet after PT 1.7 release
cc: @zcain117

[diegovalenzuelaiturra]

A version can be specified like this :

VERSION: str = "1.7"

# NOTE : use fixed source instead of "master"
#        e.g. a tag or something like "58ccf11c361e4ba2fd9d3829ca907c2d037657a1"
source: str  = "v1.7.0" 
url: str  = f"https://raw.githubusercontent.com/pytorch/xla/{source}/contrib/scripts/env-setup.py"
file: str  = "pytorch-xla-env-setup.py"

!curl $url -o $file

%run $file --version=$VERSION --apt-packages libomp5 libopenblas-dev

This may also be useful when working on TPUs on Google Colab

VERSION: str = "1.7"

wheel: str = f"torch_xla-{VERSION}-cp36-cp36m-linux_x86_64.whl"
base_url: str = "https://storage.googleapis.com/tpu-pytorch/wheels"
url: str = f"{base_url}/{wheel}"

!pip3 install cloud-tpu-client==0.10 $url