Pytorch with CUDA Support on Jetson Orin Nano: Installation and Troubleshooting

Issue Overview

Users are experiencing difficulties installing and configuring Pytorch with CUDA support on the Nvidia Jetson Orin Nano development board. The specific issue involves the inability to enable CUDA support for Pytorch, as evidenced by the torch.cuda.is_available() function returning False. This problem occurs during the setup process and affects the ability to utilize GPU acceleration for Pytorch-based applications on the Jetson Orin Nano.

The issue is observed on a system with the following specifications:

• Jetson Orin Nano
• JetPack 6
• Python 3.10
• Ubuntu 22.04 LTS
• CUDA version 12.4
• cuDNN version 8.9.4

Possible Causes

1. Incompatible Pytorch version: The initially installed Pytorch version may not be compatible with the specific CUDA version on the Jetson Orin Nano.

2. Incorrect installation process: The user might have followed a general installation guide not specific to the Jetson platform, leading to compatibility issues.

3. Missing or incompatible CUDA libraries: The CUDA toolkit installation might be incomplete or incompatible with the Pytorch version.

4. Environment configuration issues: The system environment variables might not be correctly set to point to the CUDA installation.

5. Outdated or incompatible drivers: The GPU drivers might not be up-to-date or compatible with the installed CUDA version.

Troubleshooting Steps, Solutions & Fixes

1. Use pre-built Pytorch wheels for Jetson:
   Instead of using the standard Pytorch installation, use the pre-built wheels specifically for JetPack 6 and CUDA 12.4. These can be found in the jetson-containers repository.

   Steps:
   a. Navigate to the jp6/cu124 index: https://github.com/dusty-nv/jetson-containers/tree/master/packages/pytorch/jp6/cu124
   b. Download the following wheel files:

   • torch-2.3.0
   • torchaudio-2.3.0+952ea74
   • torchvision-0.18.0a0+6043bc2
     c. Install the downloaded wheel files using pip:

   pip3 install /path/to/torch-2.3.0-wheel-file.whl
   pip3 install /path/to/torchaudio-2.3.0+952ea74-wheel-file.whl
   pip3 install /path/to/torchvision-0.18.0a0+6043bc2-wheel-file.whl
   

2. Use the l4t-pytorch container:
   If the above method doesn’t work, try using the l4t-pytorch container which has been rebuilt for CUDA 12.4, cuDNN 9, PyTorch 2.3, torchvision 0.18, and torchaudio 2.3.

   Steps:
   a. Pull the Docker image:

   docker pull dustynv/l4t-pytorch:r36.3.0-cu124
   

   b. Run the container:

   docker run --runtime nvidia -it dustynv/l4t-pytorch:r36.3.0-cu124
   

   c. Inside the container, you can test CUDA availability:

   import torch
   print(torch.cuda.is_available())
   

3. Verify CUDA installation:
   Ensure that CUDA is properly installed and recognized by the system.

   Run the following command to check CUDA version:

   nvcc --version
   

   If CUDA is not found, reinstall it using the provided commands:

   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
   

4. Check environment variables:
   Ensure that the CUDA paths are correctly set in your environment.

   Add the following lines to your ~/.bashrc file:

   export PATH=/usr/local/cuda-12.4/bin${PATH:+:${PATH}}
   export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
   

   Then, source the file:

   source ~/.bashrc
   

5. Update GPU drivers:
   Ensure that you have the latest GPU drivers installed for your Jetson Orin Nano.

   Check the current driver version:

   cat /proc/driver/nvidia/version
   

   If an update is needed, refer to the NVIDIA documentation for Jetson Orin Nano driver updates.

6. Verify Python and pip versions:
   Ensure you’re using the correct Python and pip versions.

   python3 --version
   pip3 --version
   

   If needed, upgrade pip:

   python3 -m pip install --upgrade pip
   

7. Clean installation:
   If all else fails, consider performing a clean installation of JetPack and then follow the steps to install Pytorch using the pre-built wheels or the l4t-pytorch container.

Remember to reboot your system after making significant changes to ensure all updates take effect. If the issue persists after trying these solutions, consider reaching out to the NVIDIA Jetson community forums or support channels for further assistance.