Problem

The installation of this two packages is not as straight forward as it may look. If you have already use PyTorch framework in your computer, you may be familiar with the process of installation:

pip3 install torch; pip3 install torchvision

If you install the packages using those commands, you will have success, but PyTorch will not recognize that cuda is available

import torch
torch.cuda.is_available() # This will return False

If you install packages using requirements files, you have to open them and remove both torch and torchvision, to avoid that pip installs these packages in the wrong way.

How to install PyTorch

PyTorch, although we can’t install through pip, it’s relatively easy to install. All we have to do is to follow the steps in NVIDIA documentation. Here I made a shorter step-by-step with some tricks:

• Install JetPack or check JetPack version

  • If JetPack is pre-installed, there is no direct way to check its version. You can either install a GUI version or check its L4T info:

  nano /etc/nv_tegra_release
  

  • In this file, you will find something similar to this:

  # R34 (release), REVISION: 1.1, GCID: 30414990, BOARD: t186ref, EABI: aarch64, DATE: Tue May 17 04:20:55 UTC 2022
  

  • Look for # Rxx (release) (in this case R34) and do a google search for “L4T Rxx”, and you will find documentation of your JetPack version

• Install system packages required by PyTorch (check documentation)

sudo apt-get -y update; 
sudo apt-get -y install autoconf bc build-essential g++-8 gcc-8 clang-8 lld-8 gettext-base gfortran-8 iputils-ping libbz2-dev libc++-dev libcgal-dev libffi-dev libfreetype6-dev libhdf5-dev libjpeg-dev liblzma-dev libncurses5-dev libncursesw5-dev libpng-dev libreadline-dev libssl-dev libsqlite3-dev libxml2-dev libxslt-dev locales moreutils openssl python-openssl rsync scons python3-pip libopenblas-dev;

• Enter https://developer.download.nvidia.cn/compute/redist/jp/

  • You will find a list like the one bellow:

  

  vxy is for JetPack x.y and vxyz is for JetPack x.y.z

  • Enter in your version
  • Enter in pytorch
  • Copy the link of the desired pytorch version

• Run the following (replace LinkFromPreviousStep):

export TORCH_INSTALL=**LinkFromPreviousStep**

• Then install pytorch (and its dependencies)
  • Depending on what you want, you may need to run as administrator

python3 -m pip install --no-cache $TORCH_INSTALL

Test

Open python (if you installed as administrator, now only works as administrator)

python3

Then run

import torch
torch.cuda.is_available()

If return true, the installation was successful.

Exit python with ctrl+z

Uninstall

To uninstall just use

pip3 uninstall torch

Or, if the installation was done as administrator

sudo pip3 uninstall torch