Skip to main content

Install Pytorch for reComputer Jetson

NVIDIA publishes PyTorch wheels per JetPack release, so the install differs between JetPack 5 and JetPack 6. Pick the section that matches your module, install the matching wheel, and confirm torch.cuda.is_available() returns True. Conda environments keep several torch versions side by side when a project pins an older one.

Prerequisites​

  • reComputer with JetPack installed
  • Internet connection

Install PyTorch for reComputer​

The two JetPack generations in most reComputer units take different wheels: JetPack 5.1.3 pairs with Python 3.8, JetPack 6.2 with Python 3.10.

In reComputer, the commonly used version of Jetpack 5 is 5.1.3. The steps below install PyTorch for it.

Step 1. Update system and install dependencies:

sudo apt-get -y update
sudo apt-get install -y python3-pip libopenblas-dev

Step 2. Download the official PyTorch wheel (example: torch-2.1.0 for JetPack 5.1.3, Python 3.8):

wget https://developer.download.nvidia.com/compute/redist/jp/v512/pytorch/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl
pip install torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl
note

If you need a more recent version of torch that supports Jetpack 5.1.3, please download the compiled wheel file from here (torch-2.2 for python3.8).

Step 3. Verify installation:

import torch
print(torch.cuda.is_available())
print(torch.__version__)
torch.__version__ and torch.cuda.is_available() output on JetPack 5.1.3 with torch 2.1.0

Multiple PyTorch Versions with Conda​

Miniconda is a lightweight installer for the Conda package manager, providing a minimal environment to quickly create and manage Python environments and packages.

Step 1. Install Miniconda:

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh
bash Miniconda3-latest-Linux-aarch64.sh
info
  • Follow the instructions and enter "yes" to agree to the terms.
  • The installation path is recommended to use the default setting, such as ~/miniconda3.
  • Finally, it prompts whether to automatically add to .bashrc. It is recommended to enter "yes".
conda init output showing /home/seeed/.bashrc modified so conda loads in new shells

Step 2. Initialize and configure conda:

source ~/.bashrc
# Verify if conda is functioning properly
conda --version
# Disable automatic entry into the conda base environment
conda config --set auto_activate_base false
# Close and reopen the terminal. By default, it will not enter the base environment.

Step 3. Create a new environment and install a specific PyTorch wheel:

conda create -n torch_2.0 python=3.8
conda activate torch_2.0

Step 4. Download and install the PyTorch wheel file as shown in the previous sections.

info

For more details, see the NVIDIA PyTorch installation guide

Resources​

FAQ​

Which PyTorch wheel matches my JetPack version? JetPack 5.1.3 uses torch-2.1.0a0+...cp38 (Python 3.8) and JetPack 6.2 uses torch-2.5.0a0+...cp310 (Python 3.10). Pick the wheel whose jp/vXX path and cpXX tag match both, otherwise pip installs something that runs on CPU only.

Do I need cuSPARSELt before installing torch? Only on the newer wheels. PyTorch builds from 24.06 onward call cuSPARSELt for sparse matmul, so the JetPack 6.2 path installs 0.7.1 first. The JetPack 5.1.3 example with torch 2.1.0 does not need it.

How do I keep two torch versions on one machine? Use conda environments: one env per Python and torch pair, for example conda create -n torch_2.0 python=3.8. Activate the env before installing the wheel, and keep auto_activate_base off so the system Python stays clean.

torch.cuda.is_available() returns False. What went wrong? The wheel is either the CPU build from PyPI or built for a different JetPack. Reinstall the NVIDIA wheel that matches your JetPack path and confirm with torch.__version__, which carries the nv build tag for CUDA-enabled wheels.

Tech Support & Product Discussion​

Thank you for choosing our products! We are here to provide you with different support to ensure that your experience with our products is as smooth as possible. We offer several communication channels to cater to different preferences and needs.

Loading Comments...