Pytorch for JP 6.2 running cuda 12.6

[charan.kumar]

Hi,

I have been trying to install pytorch in my Jetson AGX Orin Dev kit, but I cannot find the pytorch wheel file that would run with cuda 12.6. is there a prebuilt wheel for JP6.2, CUDA 12.6? if not what are my other options?

[kesong]

You can have a try the pytorch wheel here: jp6/cu126 index

[charan.kumar]

The Torch wheel available at jp6/cu126 index is built with NumPy version 1.x, while other components of my project—such as the ZED camera package—require NumPy 2.2.6. Because of this incompatibility, I can’t use the provided Torch wheel. Since there are no compatible Torch wheels available for JetPack 6.2, do I need to downgrade to JetPack 6.0?

[whitesscott]

#Maybe first check ultralytics and see what version of torch it wants.

sudo apt update && sudo apt install python3.10-venv #or whatever version python3 you are using.

python3 -m venv torchenv
source torchenv/bin/activate

pip install numpy ninja cmake

wget https://github.com/pytorch/pytorch/releases/download/v2.7.0/pytorch-v2.7.0.tar.gz
tar xfz pytorch-v2.7.0.tar.gz
cd pytorch

#look at setup.py for other environment variables you can set if you need them.
export USE_ROCM=0
export MAX_JOBS=6
export TORCH_CUDA_ARCH_LIST=“8.7”
export USE_CUDA=ON
export USE_PRIORITIZED_TEXT_FOR_LD=1
export LD_LIBRARY_PATH=‘/usr/local/cuda/lib64:$LD_LIBRARY_PATH’

#this takes a while to compile.
python3 setup.py bdist_wheel --dist-dir dist

#only if after some minutes it fails do the following 4 lines not counting the comment.
cd build
ninja -j 6
#when that completes run this again. It will quickly build the torch*.whl
cd ..
python setup.py bdist_wheel --dist-dir dist

#then after the wheel is built
deactivate #to exit venv

#then
pip install dist/torch*.whl