PyTorch for Jetson Platform

NVIDIA Optimized Frameworks (Latest Release) Download 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.4.0a0+f70bd71a48 24.06 24.06 6.0
2.4.0a0+07cecf4168 24.05 24.05 6.0
2.3.0a0+6ddf5cf85e 24.04 24.04 6.0 Developer Preview
2.3.0a0+40ec155e58 24.03 24.03
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

  • TensorRT Model Optimizer is a new module but is not yet supported on Jetson due to its dependency on pytorch distributed functionality, which was disabled for Jetson.
© Copyright 2024, NVIDIA. Last updated on Jun 27, 2024.