TPU setup in docs results in import failure for flash.Task

[fonnesbeck]

I am following TPU setup instructions on a clean Colab notebook. When I do so and attempt to import flash.Task I get the following error:

ImportError: /usr/local/lib/python3.7/dist-packages/_XLAC.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZN2at13_foreach_add_EN3c108ArrayRefINS_6TensorEEENS1_INS0_6ScalarEEE

This is being run on Colab Pro with TPU enabled.

[fonnesbeck]

Actually, looking at the install log, I see the following:

ERROR: torch-1.15-cp36-cp36m-linux_x86_64.whl is not a supported wheel on this platform.
ERROR: torch_xla-1.15-cp36-cp36m-linux_x86_64.whl is not a supported wheel on this platform.
ERROR: torchvision-1.15-cp36-cp36m-linux_x86_64.whl is not a supported wheel on this platform.

It looks like the install instructions need to be updated.

[awaelchli]

Here a Google Colab that shows the issue:
https://colab.research.google.com/github/PytorchLightning/pytorch-lightning/blob/master/notebooks/06-mnist-tpu-training.ipynb
need to find new wheels and then update the notebooks

[awaelchli]

oh wow
it's because google colab recently switched the default python version from python 3.6 to 3.7.
This link works:
! pip install https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.7-cp37-cp37m-linux_x86_64.whl

You just have to replace 36 with 37 :)
I will update the notebooks