Using Pytorch on GPU after upgrading cuda to 11.8 on Jetpack5

[dusty_nv]

@YoushaaMurhij you might want to try re-installing the PyTorch 2.0 wheel that has CUDA support from here:

PyTorch for Jetson

Below are pre-built PyTorch pip wheel installers for Jetson Nano, TX1/TX2, Xavier, and Orin with JetPack 4.2 and newer. Download one of the PyTorch binaries from below for your version of JetPack, and see the installation instructions to run on your Jetson. These pip wheels are built for ARM aarch64 architecture, so run these commands on your Jetson (not on a host PC). You can also use the pre-built l4t-pytorch and l4t-ml container images and Dockerfiles. PyTorch pip wheels JetPack 5

PyTo…

Or baring that, rebuild PyTorch from source for CUDA 11.8 following the Build from Source instructions from that thread. I haven’t tried upgrading CUDA with PyTorch before, so YMMV. Also if/when you go to build torchvision, the CUDA versions will need to match otherwise PyTorch will throw an error about it.