PyTorch for JetPack 4.4 - L4T R32.4.3 in Jetson Xavier NX

NOTE: The JetPack 4.4 production release (L4T R32.4.3) only supports PyTorch 1.6.0 or newer, due to updates in cuDNN.

Installation for Python 3.6

  1. PyTorch 1.6.0

wget https://nvidia.box.com/shared/static/yr6sjswn25z7oankw8zy1roow9cy5ur1.whl -O torch-1.6.0rc2-cp36-cp36m-linux_aarch64.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev
sudo pip3 install Cython
sudo pip3 install torch-1.6.0rc2-cp36-cp36m-linux_aarch64.whl

  1. Torchvision 0.6.0

sudo apt-get install libjpeg-dev zlib1g-dev
git clone --branch v0.6.0 https://github.com/pytorch/vision torchvision
cd torchvision
sudo python3 setup.py install

Test

  1. Pytourch
import torch
print(torch.__version__)
print('CUDA available: ' + str(torch.cuda.is_available()))
print('cuDNN version: ' + str(torch.backends.cudnn.version()))
a = torch.cuda.FloatTensor(2).zero_()
print('Tensor a = ' + str(a))
b = torch.randn(2).cuda()
print('Tensor b = ' + str(b))
c = a + b
print('Tensor c = ' + str(c))
  1. Torchvision

import torchvision
print(torchvision.version)

Documentation - PyTorch for Jetson

2 Likes

That’s correct, thanks for the help to clarify this.

1 Like

Those directions worked perfectly for me. Thanks!