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?