Pytorch with CUDA support Jetson Orin Nano

[leetianhau]

Hi everyone,
I need some help regarding building Pytorch with Cuda on my Jetson Orin Nano which is using Jetpack 6, Python 3.10, Ubuntu 22.04 LTS ,a CUDA version of 12.4 (from jtop) abd cuDNN version of 8.9.4 (from jtop).

From what I did, I went to download CUDA Toolkit 12.4 and Pytorch with CUDA 12.1 support. Then I open my Python 3.10 to run torch.cuda.is_available() and it said false.

CUDA Toolkit 12.4 : CUDA Toolkit 12.4 Update 1 Downloads | NVIDIA Developer
codes given from website above:
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/arm64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-4 cuda-compat-12-4

Pytorch with CUDA support website: Start Locally | PyTorch
codes:
pip3 install torch torchvision torchaudio

I have no idea at this point, is there someone who can show me the steps to do it?

[dusty_nv]

Hi @leetianhau, you can find the PyTorch wheels for JetPack 6 and CUDA 12.4 built from jetson-containers/packages/pytorch at master · dusty-nv/jetson-containers · GitHub here: jp6/cu124 index

[leetianhau]

Hi @dusty_nv , thank you so much for replying. So I just have to download [torch-2.3.0], [torchaudio-2.3.0+952ea74] and [torchvision-0.18.0a0+6043bc2] ??
I downloaded these three and run pip3 install to these three filenames but when I run
torch.cuda.is_available() and it gave me False.

[dusty_nv]

Hi @leetianhau, I rebuilt l4t-pytorch container for CUDA 12.4, cuDNN 9, PyTorch 2.3, torchvision 0.18, and torchaudio 2.3. It passes the tests here using GPU. Can you give it a try?

• dustynv/l4t-pytorch:r36.3.0-cu124
• jetson-containers/packages/l4t/l4t-pytorch at master · dusty-nv/jetson-containers · GitHub