jetson-containers/packages/pytorch/version.py at master · KTH-SML/jetson-containers · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# this is in a different file from config.py so other packages can import it
from jetson_containers import L4T_VERSION, CUDA_VERSION, SYSTEM_ARM
from packaging.version import Version

import os

if 'PYTORCH_VERSION' in os.environ and len(os.environ['PYTORCH_VERSION']) > 0:
    PYTORCH_VERSION = Version(os.environ['PYTORCH_VERSION'])
elif SYSTEM_ARM:
    if L4T_VERSION.major >= 38:
        if CUDA_VERSION >= Version('13.0'):   # JetPack 7 (CUDA 13.0)
            PYTORCH_VERSION = Version('2.8')
        elif CUDA_VERSION >= Version('12.9'):   # JetPack 7.0 (CUDA 12.9)
            PYTORCH_VERSION = Version('2.8')
        else:
            PYTORCH_VERSION = Version('2.8')  # JetPack 7.0 (CUDA 12.9)
    elif L4T_VERSION.major >= 36:
        if CUDA_VERSION >= Version('13.0'):   # JetPack 6.2 (CUDA 12.6)
            PYTORCH_VERSION = Version('2.8')
        elif CUDA_VERSION >= Version('12.9'):   # JetPack 6.2 (CUDA 12.6)
            PYTORCH_VERSION = Version('2.8')
        elif CUDA_VERSION >= Version('12.8'):   # JetPack 6.2 (CUDA 12.6)
            PYTORCH_VERSION = Version('2.7')
        elif CUDA_VERSION == Version('12.6'):   # JetPack 6.2 (CUDA 12.6)
            PYTORCH_VERSION = Version('2.8')
        elif CUDA_VERSION >= Version('12.4'): # JetPack 6.0 (CUDA 12.4)
            PYTORCH_VERSION = Version('2.4')
        else:
            PYTORCH_VERSION = Version('2.2')  # JetPack 6.0 (CUDA 12.2)
    elif L4T_VERSION.major >= 34:
        PYTORCH_VERSION = Version('2.2')      # JetPack 5.1 (CUDA 11.4)
    elif L4T_VERSION.major >= 32:
        PYTORCH_VERSION = Version('1.10')     # JetPack 4.6 (CUDA 10.2)
else:
    PYTORCH_VERSION = Version('2.8')          # pytorch nightly (CUDA 13.0)