Installing Torch/Torchvision on Jetson Orin Nano Jetpack 6.2.2 Cuda 12.6

After a bit of time attempting to install torch and torchvision on the Jetson Orin Nano I have been able to get this working.

This guide is provided for educational purposes only. I’m not responsible for any damage, data loss, or malfunction caused by following these instructions. Use at your own risk.

System Specs:

Kernel name: Linux ubuntu 5.15.185-tegra

Ubuntu Release: 22.04.5 LTS (Jammy Jellyfish)

Jetpack: 6.2.2

Cuda: 12.6

Model: Jetson Orin Nano 8GB Devkit

##############################################################################

INSTALL cusparselt for cuda-12:

Set target platform and download .deb:

TAKEN FROM DOCS MAY CHANGE IN TIME:

wget https://developer.download.nvidia.com/compute/cusparselt/0.8.1/local_installers/cusparselt-local-tegra-repo-ubuntu2204-0.8.1_0.8.1-1_arm64.deb
sudo dpkg -i cusparselt-local-tegra-repo-ubuntu2204-0.8.1_0.8.1-1_arm64.deb
sudo cp /var/cusparselt-local-tegra-repo-ubuntu2204-0.8.1/cusparselt-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cusparselt
sudo apt-get -y install cusparselt-cuda-12

DOWNGRADE numpy NUMPY VERSION 2 NOT SUPPORTED WHILE CURRENTLY TESTING:

pip install “numpy<2” --force-reinstall

INSTALL torch

pip install https://developer.download.nvidia.com/compute/redist/jp/v61/pytorch/torch-2.5.0a0+872d972e41.nv24.08.17622132-cp310-cp310-linux_aarch64.whl

CHECK IF TORCH IS WORKING AND USING GPU SHOULD RETURN ‘TRUE’

python3 -c “import torch; print(‘torch:’, torch.version, ‘CUDA:’, torch.version.cuda, ‘GPU available:’, torch.cuda.is_available())”

BUILD torchvision FROM SOURCE:

git clone GitHub - pytorch/vision: Datasets, Transforms and Models specific to Computer Vision · GitHub
cd vision
git checkout v0.20.0

SET THESE FLAGS

export BUILD_VERSION=0.20.0
export TORCH_CUDA_ARCH_LIST=“8.7”
export FORCE_CUDA=1
export CUDA_HOME=/usr/local/cuda
export MAX_JOBS=2

INSTALL DEPENDENCIES

sudo apt install -y libjpeg-dev libpng-dev libavcodec-dev libavformat-dev libswscale-dev

python3 setup.py install

AFTER COMPILATION

RUN cd “change to another DIR”

RUN TO CHECK torchvision is WORKING AND USING GPU:

"
python3 - <<EOF
import torch
import torchvision
from torchvision.ops import nms

print(“torch:”, torch.version)
print(“torchvision:”, torchvision.version)
print(“cuda:”, torch.cuda.is_available())

boxes = torch.tensor([[0,0,10,10],[1,1,11,11]], dtype=torch.float32).cuda()
scores = torch.tensor([0.9, 0.8]).cuda()

print(“nms works:”, nms(boxes, scores, 0.5))
EOF
"

SHOULD DISPLAY

torch: 2.5.0a0+872d972e41.nv24.08
torchvision: 0.20.0
cuda: True
nms works: tensor([0], device=‘cuda:0’)

OPTIONAL INSTALL opencv and ultranalytics “keeping numpy version”:

pip install “numpy<2” ultralytics

pip install --force-reinstall --no-deps opencv-python

##############################################################################

Hi,

Just FYI.
You can also find the package built for JetPack 6 in the apt server below:

Thanks.