Can not import torch_xla on google colab without tpu

Hi,
I want to use torch_xla to verify some problems on Google Colab. Due to the lack of TPU resources, I am using a CPU environment. The default environment for Colab is Python 3.9, but the official tutorial for torch_xla uses 3.8. Therefore, I switch Python to 3.8 and installed torch and torch-xla according to the tutorial. I can see the torch-xla module when I use the ‘pip3 list’ command. However, I encountered an error when importing torch_xla: ModuleNotFoundError: No module named ‘torch_xla’
I would like to know if torch_xla must run in a TPU environment. If not, where did I go wrong? Thanks a lot.
Here is my Google Colab ipynb file

command and output

!sudo apt-get install python3.8 python3.8-dev python3.8-distutils python3.8-venv

Reading package lists… Done
Building dependency tree
Reading state information… Done
Note, selecting ‘python3-distutils’ instead of ‘python3.8-distutils’
python3-distutils is already the newest version (3.8.10-0ubuntu1~20.04).
python3-distutils set to manually installed.
python3.8 is already the newest version (3.8.10-0ubuntu1~20.04.6).
python3.8 set to manually installed.
python3.8-dev is already the newest version (3.8.10-0ubuntu1~20.04.6).
python3.8-dev set to manually installed.
The following additional packages will be installed:
python-pip-whl
The following NEW packages will be installed:
python-pip-whl python3.8-venv
0 upgraded, 2 newly installed, 0 to remove and 22 not upgraded.
Need to get 1,810 kB of archives.
After this operation, 2,339 kB of additional disk space will be used.
Get:1 amd64 python-pip-whl all 20.0.2-5ubuntu1.8 [1,805 kB]
Get:2 amd64 python3.8-venv amd64 3.8.10-0ubuntu1~20.04.6 [5,448 B]
Fetched 1,810 kB in 1s (2,609 kB/s)
debconf: unable to initialize frontend: Dialog
debconf: (No usable dialog-like program is installed, so the dialog based frontend cannot be used. at /usr/share/perl5/Debconf/FrontEnd/Dialog.pm line 76, <> line 2.)
debconf: falling back to frontend: Readline
debconf: unable to initialize frontend: Readline
debconf: (This frontend requires a controlling tty.)
debconf: falling back to frontend: Teletype
dpkg-preconfigure: unable to re-open stdin:
Selecting previously unselected package python-pip-whl.
(Reading database … 128275 files and directories currently installed.)
Preparing to unpack …/python-pip-whl_20.0.2-5ubuntu1.8_all.deb …
Unpacking python-pip-whl (20.0.2-5ubuntu1.8) …
Selecting previously unselected package python3.8-venv.
Preparing to unpack …/python3.8-venv_3.8.10-0ubuntu1~20.04.6_amd64.deb …
Unpacking python3.8-venv (3.8.10-0ubuntu1~20.04.6) …
Setting up python-pip-whl (20.0.2-5ubuntu1.8) …
Setting up python3.8-venv (3.8.10-0ubuntu1~20.04.6) …

!sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.8 1

!sudo update-alternatives --config python3

There are 2 choices for the alternative python3 (providing /usr/bin/python3).

Selection Path Priority Status

  • 0 /usr/bin/python3.9 2 auto mode
    1 /usr/bin/python3.8 1 manual mode
    2 /usr/bin/python3.9 2 manual mode

Press to keep the current choice[*], or type selection number: 1
update-alternatives: using /usr/bin/python3.8 to provide /usr/bin/python3 (python3) in manual mode

!python3 --version

Python 3.8.10

!sudo apt install python3-pip

Reading package lists… Done
Building dependency tree
Reading state information… Done
The following additional packages will be installed:
python3-setuptools python3-wheel
Suggested packages:
python-setuptools-doc
The following NEW packages will be installed:
python3-pip python3-setuptools python3-wheel
0 upgraded, 3 newly installed, 0 to remove and 22 not upgraded.
Need to get 584 kB of archives.
After this operation, 2,621 kB of additional disk space will be used.
Get:1 amd64 python3-setuptools all 45.2.0-1ubuntu0.1 [330 kB]
Get:2 amd64 python3-wheel all 0.34.2-1ubuntu0.1 [23.9 kB]
Get:3 amd64 python3-pip all 20.0.2-5ubuntu1.8 [231 kB]
Fetched 584 kB in 1s (1,022 kB/s)
debconf: unable to initialize frontend: Dialog
debconf: (No usable dialog-like program is installed, so the dialog based frontend cannot be used. at /usr/share/perl5/Debconf/FrontEnd/Dialog.pm line 76, <> line 3.)
debconf: falling back to frontend: Readline
debconf: unable to initialize frontend: Readline
debconf: (This frontend requires a controlling tty.)
debconf: falling back to frontend: Teletype
dpkg-preconfigure: unable to re-open stdin:
Selecting previously unselected package python3-setuptools.
(Reading database … 128312 files and directories currently installed.)
Preparing to unpack …/python3-setuptools_45.2.0-1ubuntu0.1_all.deb …
Unpacking python3-setuptools (45.2.0-1ubuntu0.1) …
Selecting previously unselected package python3-wheel.
Preparing to unpack …/python3-wheel_0.34.2-1ubuntu0.1_all.deb …
Unpacking python3-wheel (0.34.2-1ubuntu0.1) …
Selecting previously unselected package python3-pip.
Preparing to unpack …/python3-pip_20.0.2-5ubuntu1.8_all.deb …
Unpacking python3-pip (20.0.2-5ubuntu1.8) …
Setting up python3-setuptools (45.2.0-1ubuntu0.1) …
Setting up python3-wheel (0.34.2-1ubuntu0.1) …
Setting up python3-pip (20.0.2-5ubuntu1.8) …
Processing triggers for man-db (2.9.1-1) …

!pip3 install cloud-tpu-client==0.10 torch==1.13.0 https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-1.13-cp38-cp38-linux_x86_64.whl

Looking in indexes: ://pypi.org/simple, ://us-python.pkg.dev/colab-wheels/public/simple/
Collecting cloud-tpu-client==0.10
Using cached cloud_tpu_client-0.10-py3-none-any.whl (7.4 kB)
Collecting torch==1.13.0
Using cached torch-1.13.0-cp38-cp38-manylinux1_x86_64.whl (890.2 MB)
Collecting torch-xla==1.13
Using cached ://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-1.13-cp38-cp38-linux_x86_64.whl (151.3 MB)
Collecting google-api-python-client==1.8.0
Using cached google_api_python_client-1.8.0-py3-none-any.whl (57 kB)
Collecting oauth2client
Using cached oauth2client-4.1.3-py2.py3-none-any.whl (98 kB)
Collecting typing-extensions
Downloading typing_extensions-4.5.0-py3-none-any.whl (27 kB)
Collecting nvidia-cudnn-cu11==8.5.0.96
Downloading nvidia_cudnn_cu11-8.5.0.96-2-py3-none-manylinux1_x86_64.whl (557.1 MB)
anylinux1_x86_64.whl (21.0 MB)
anylinux1_x86_64.whl (317.1 MB)
e-cu11==11.7.99
Downloading nvidia_cuda_runtime_cu11-11.7.99-py3-none-manylinux1_x86_64.whl (849 kB)
plate<4dev,>=3.0.0
Downloading uritemplate-3.0.1-py2.py3-none-any.whl (15 kB)
Collecting google-auth-httplib2>=0.0.3
Downloading google_auth_httplib2-0.1.0-py2.py3-none-any.whl (9.3 kB)
Collecting google-auth>=1.4.1
Downloading google_auth-2.16.2-py2.py3-none-any.whl (177 kB)
odules>=0.0.5
Downloading pyasn1_modules-0.2.8-py2.py3-none-any.whl (155 kB)
ent already satisfied: wheel in /usr/lib/python3/dist-packages (from nvidia-cublas-cu11==11.10.3.66->torch==1.13.0) (0.34.2)
Requirement already satisfied: setuptools in /usr/lib/python3/dist-packages (from nvidia-cublas-cu11==11.10.3.66->torch==1.13.0) (45.2.0)
Collecting protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<4.0.0dev,>=3.19.5
Downloading protobuf-3.20.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (1.0 MB)
mon-protos<2.0dev,>=1.56.2
Downloading googleapis_common_protos-1.58.0-py2.py3-none-any.whl (223 kB)
alizer<4,>=2
Downloading charset_normalizer-3.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (195 kB)
odules, google-auth, idna, charset-normalizer, urllib3, certifi, requests, googleapis-common-protos, google-api-core, pyparsing, httplib2, uritemplate, google-auth-httplib2, google-api-python-client, oauth2client, cloud-tpu-client, typing-extensions, nvidia-cublas-cu11, nvidia-cudnn-cu11, nvidia-cuda-nvrtc-cu11, nvidia-cuda-runtime-cu11, torch, absl-py, torch-xla
Successfully installed absl-py-1.4.0 cachetools-5.3.0 certifi-2022.12.7 charset-normalizer-3.1.0 cloud-tpu-client-0.10 google-api-core-1.34.0 google-api-python-client-1.8.0 google-auth-2.16.2 google-auth-httplib2-0.1.0 googleapis-common-protos-1.58.0 httplib2-0.21.0 idna-3.4 nvidia-cublas-cu11-11.10.3.66 nvidia-cuda-nvrtc-cu11-11.7.99 nvidia-cuda-runtime-cu11-11.7.99 nvidia-cudnn-cu11-8.5.0.96 oauth2client-4.1.3 protobuf-3.20.3 pyasn1-0.4.8 pyasn1-modules-0.2.8 pyparsing-3.0.9 requests-2.28.2 rsa-4.9 six-1.16.0 torch-1.13.0 torch-xla-1.13 typing-extensions-4.5.0 uritemplate-3.0.1 urllib3-1.26.15

{“pip_warning”:{“packages”:[“certifi”,“google”]}}

import torch

!pip3 list

Package Version


absl-py 1.4.0
cachetools 5.3.0
certifi 2022.12.7
charset-normalizer 3.1.0
cloud-tpu-client 0.10
dbus-python 1.2.16
google-api-core 1.34.0
google-api-python-client 1.8.0
google-auth 2.16.2
google-auth-httplib2 0.1.0
googleapis-common-protos 1.58.0
httplib2 0.21.0
idna 3.4
nvidia-cublas-cu11 11.10.3.66
nvidia-cuda-nvrtc-cu11 11.7.99
nvidia-cuda-runtime-cu11 11.7.99
nvidia-cudnn-cu11 8.5.0.96
oauth2client 4.1.3
pip 20.0.2
protobuf 3.20.3
pyasn1 0.4.8
pyasn1-modules 0.2.8
PyGObject 3.36.0
pyparsing 3.0.9
python-apt 2.0.1
requests 2.28.2
requests-unixsocket 0.2.0
rsa 4.9
screen-resolution-extra 0.0.0
setuptools 45.2.0
six 1.16.0
torch 1.13.0
torch-xla 1.13
typing-extensions 4.5.0
uritemplate 3.0.1
urllib3 1.26.15
wheel 0.34.2
xkit 0.0.0

import torch_xla


ModuleNotFoundError Traceback (most recent call last)
in
----> 1 import torch_xla

ModuleNotFoundError: No module named ‘torch_xla’


NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
“Open Examples” button below.

Hi, sorry for the late reply. XLA can be run with CPU, GPU and other accelerators. I suspect what you encountered is a pip issue. We shipped the python 3.9 whl for the 2.0 release. For the colab you can use https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-2.0-cp39-cp39-linux_x86_64.whl. A complete list of whls is under GitHub - pytorch/xla: Enabling PyTorch on Google TPU.

Btw we check github issue more often, feel free to open a github issue under Issues · pytorch/xla · GitHub if you run into any additional issues.

BTW, we check