Jetson Orin NX (JetPack 5.1.5 with cuda11.8) how to install pytorch

Hello everyone,

I’m currently running into issues installing PyTorch and its associated packages on my Jetson Orin NX (JetPack 5.1.5), and I’d really appreciate some guidance here.

Environment

  • Device: Jetson Orin NX

  • JetPack version: 5.1.5 (default CUDA version: 11.4; I’m targeting CUDA 11.8)

  • Python setup: Miniconda3 virtual environment (vings_vio) with Python 3.9.19

Required Packages

I need to install these specific versions:

  • torch==2.0.1+cu118

  • torch-scatter==2.0.2

  • torchaudio==2.0.2+cu118

  • torchmetrics==1.4.0

  • torchvision==0.15.2+cu118

Steps I Took

  1. Replaced CUDA 11.4 with CUDA 11.8 using these commands:
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/arm64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-tegra-repo-ubuntu2004-11-8-local_11.8.0-1_arm64.deb
sudo dpkg -i cuda-tegra-repo-ubuntu2004-11-8-local_11.8.0-1_arm64.deb
sudo cp /var/cuda-tegra-repo-ubuntu2004-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

  1. Created the Miniconda3 virtual environment vings_vio (Python 3.9.19).

  2. Attempted to install the required packages in the virtual environment.

Issues Encountered

  • When trying to install torch==2.0.1+cu118, I got:
Looking in indexes: https://download.pytorch.org/whl/cu118
ERROR: Could not find a version that satisfies the requirement torch==2.0.1 (from versions: none)
ERROR: No matching distribution found for torch==2.0.1

  • When installing torch-scatter==2.0.2, I hit a ModuleNotFoundError:
Looking in links: https://data.pyg.org/whl/torch-2.0.2+cu118.html
Collecting torch-scatter==2.0.2
  Using cached torch_scatter-2.0.2.tar.gz (22 kB)
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error
  
  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [15 lines of output]
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
          exec(compile('''
          ~~~~^^^^^^^^^^^^
          # This is <pip-setuptools-caller> -- a caller that pip uses to run setup.py
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
          ...<32 lines>...
          exec(compile(setup_py_code, filename, "exec"))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
          ''' % ('/tmp/pip-install-_ht8_397/torch-scatter_b6ba230c5a2b41a9afe0411f0cdc49c3/setup.py',), "<pip-setuptools-caller>", "exec"))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "<pip-setuptools-caller>", line 35, in <module>
        File "/tmp/pip-install-_ht8_397/torch-scatter_b6ba230c5a2b41a9afe0411f0cdc49c3/setup.py", line 7, in <module>
          import torch
      ModuleNotFoundError: No module named 'torch'
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

  • I also got:
ERROR: Directory 'submodules/diff-surfel-rasterization' is not installable. Neither 'setup.py' nor 'pyproject.toml' found.

Current Status

  • I tried installing PyTorch via this command:
pip3 install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118

  1. pip3 list confirms these are installed:
  • torch                         2.0.1
    torchaudio                    2.0.2
    torchvision                   0.15.2
    
    
  • But CUDA isn’t detected when verifying:

(vings_vio) mickrobot@ubuntu:~/VINGS-Mono$ python3 -c "import torch; print('Torch版本:', torch.__version__); print('CUDA是否可用:', torch.cuda.is_available())"

Output:

Torch版本: 2.0.1
CUDA是否可用: False

Could anyone help me properly install these packages (with CUDA 11.8 support) on Jetson Orin NX (JetPack 5.1.5)? I mainly need to get CUDA working with PyTorch 2.0.1, and install torch-scatter & other required packages correctly.

When I verified CUDA installation: Ran nvcc -V in the vings_vio virtual environment, output:

(vings_vio) mickrobot@ubuntu:~/VINGS-Mono/temp$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:43:33_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0

Hi,

Please visit the PyTorch for Jetson page to download the prebuilt PyTorch wheels compatible with JetPack 5.x for Jetson devices.

The commands you mentioned are for desktop gpu.

Thanks