Unable to install pytorch with cuda support on jetson AGX Xavier

I want to install pytorch with cuda support on jetson AGX Xavier (L4T 35.3.1).

From the https://pytorch.org/ , I came to know that pytorch only supports specific versions of cuda . In my case pytorch supports cuda 11.8 and 12.1 . But the one(cuda) which is default on L4T 35.3.1 is 12.2 . Anyhow I managed to uninstall it and installed the 11.8 version CUDA Toolkit 11.8 Downloads | NVIDIA Developer . Now my /usr/local/cuda/bin/nvcc --version shows 11.8.

Now I install pytorch with “pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118” command.

After that I import torch and check for torch.cuda.is_available() . It still says false. How do I enable torch build with cuda on jetson? Any suggestions are welcomed.

Additional hint : on page http://download.pytorch.org/whl/cu118/torch/ page , there are many wheel file . For amd64 and x86 , the package name contains cu118(torch-2.0.0+cu118-cp38-cp38-linux_x86_64.whl ) . But for arch64 the package name doesnt contain cu118 ( torch-1.12.0-cp38-cp38-manylinux2014_aarch64.whl ) . I tried installing torch-1.12.0-cp38-cp38-manylinux2014_aarch64.whl , but torch.cuda.is_available() is still false.

@vikaash.kb revert to the original version of CUDA for JetPack 5 (which is CUDA 11.4) and then install one of the PyTorch wheels built for JetPack 5 from here:

You can also find instructions for building PyTorch from source in that thread if you desire.

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.