PyTorch for Jetson Platform

NVIDIA Docs Hub NVIDIA Optimized Frameworks NVIDIA Optimized Frameworks PyTorch for Jetson Platform

PyTorch for Jetson Platform (PDF)

This document describes the key features, software enhancements and improvements, and known issues regarding PyTorch on the Jetson platform. See Installing PyTorch for Jetson Platform for installation information.

Key Features and Enhancements

This release includes the following key features and enhancements.

The TF32 numerical format is enabled by default for cuBLAS and cuDNN operations on Ampere GPUs starting with the 22.06 release. If you encounter training issues especially for regression, generative or higher-order models, or by using TF32 operations in pre- or post-processing steps, try to disable TF32 by setting the following:
Copy

Copied!
```
            
            torch.set_float32_matmul_precision('highest')
        
```

Compatibility

Table 1. PyTorch compatibility with NVIDIA containers and Jetpack
PyTorch Version	NVIDIA Framework Container	NVIDIA Framework Wheel	JetPack Version
2.3.0a0+40ec155e58	24.03	24.03	6.0 Developer Preview
2.3.0a0+ebedce2	24.02	24.02
2.2.0a0+81ea7a4	23.12, 24.01	23.12, 24.01
2.2.0a0+6a974bec	23.11	23.11
2.1.0a		23.06	5.1.x
2.0.0		23.05
2.0.0a0+fe05266f		23.04
2.0.0a0+8aa34602		23.03
1.14.0a0+44dac51c		23.02, 23.01
1.13.0a0+936e930		22.11	5.0.2
1.13.0a0+d0d6b1f		22.09, 22,10
1.13.0a0+08820cb	22.07	22.07
1.13.0a0+340c412	22.06	22.06	5.0.1
1.12.0a0+8a1a93a9	22.05	22.05	5.0
1.12.0a0+bd13bc66		22.04
1.12.0a0+2c916ef		22.03
1.11.0a0+bfe5ad28		22.01	4.6.1

Using PyTorch with the Jetson Platform

Storage: If you need more storage, we recommend connecting an external SSD via SATA on TX2 or Xavier devices, or USB on Jetson Nano.

Known Issues

There is a known CUPTI permissions issue that has been there since the 24.01 release. This issue prevents the profiler from being able to capture cuda events and manifests itself as a CUPTI Runtime Error with error code 35. This may be worked around by running the following command inside the container: rm -rf /usr/local/cuda/compat/lib.real. This issue would be fixed in the 24.03 release.There is a known CUPTI permissions issue that has been there since the 24.01 release. This issue prevents the profiler from being able to capture cuda events and manifests itself as a CUPTI Runtime Error with error code 35. This may be worked around by running the following command inside the container: rm -rf /usr/local/cuda/compat/lib.real.
AttributeError: module 'torch.distributed' has no attribute '_all_gather_base' would be encountered if users are importing apex in the code.