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

print(cuda.__version__)
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.current_device())

produces:

12.0.0
1.13.1
False
Traceback (most recent call last):
  File "/home/nick/yolo/cuda-2.py", line 6, in <module>
    print(torch.cuda.current_device())
  File "/home/nick/yolo/yolo-pip/lib/python3.8/site-packages/torch/cuda/__init__.py", line 552, in current_device
    _lazy_init()
  File "/home/nick/yolo/yolo-pip/lib/python3.8/site-packages/torch/cuda/__init__.py", 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?