Need Torch compiled with CUDA enabled

I have CUDA 12.0 installed on Jetson AGX Orin (Aarch64). Following instruction on Start Locally | PyTorch, I run:

pip3 install cuda-python
pip3 install torch torchvision torchaudio

I expected to have Torch compiled with CUDA, but this script:

import cuda, torch



Traceback (most recent call last):
  File "/home/nick/yolo/", line 6, in <module>
  File "/home/nick/yolo/yolo-pip/lib/python3.8/site-packages/torch/cuda/", line 552, in current_device
  File "/home/nick/yolo/yolo-pip/lib/python3.8/site-packages/torch/cuda/", line 221, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

which contradicts the installation instructions. How can I get Torch with CUDA enabled?

The built binaries mentioned on the PyTorch website do not support Jetson Orin and you should use these instructions.

1 Like

Thank you. I will give it a try. Is there a note about no support for Jetson Orin on PyTorch page?

It works with some caveats, see: Wrong instructions in Installing PyTorch for Jetson Platform - Jetson AGX Orin - NVIDIA Developer Forums.

Thanks for letting me know you are seeing issues with the install instructions. We’ll take a look at it.

Thank you. Can you take a look at a related question about Torchvision: Is there a Pip installable version of TorchVision for Jetson AGX Orin? - Jetson AGX Orin - NVIDIA Developer Forums?