Add PyTorch 2.6.0 support

MiroPsota · Feb 3, 2025 · 5dd8066 · 5dd8066
1 parent 3036ef3
commit 5dd8066
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 1 deletion.
diff --git a/.github/workflows/cuda/Linux_env.sh b/.github/workflows/cuda/Linux_env.sh
@@ -8,6 +8,7 @@ if [[ $CUDA_VERSION == "cu102" ]]; then
 fi
 
 # https://github.com/pytorch/builder/blob/main/manywheel/build_cuda.sh
+# https://github.com/pytorch/pytorch/blob/main/.ci/manywheel/build_cuda.sh
 declare -A CUDA_ARCHS=(
   ["cu102"]="3.7;5.0;6.0;7.0"
 
@@ -19,6 +20,7 @@ declare -A CUDA_ARCHS=(
 
   ["cu121"]="5.0;6.0;7.0;7.5;8.0;8.6;9.0"
   ["cu124"]="5.0;6.0;7.0;7.5;8.0;8.6;9.0"
+  ["cu126"]="5.0;6.0;7.0;7.5;8.0;8.6;9.0+PTX"
 )
 export TORCH_CUDA_ARCH_LIST=${CUDA_ARCHS[$CUDA_VERSION]}
 export FORCE_CUDA=1
diff --git a/.github/workflows/cuda/Windows_env.sh b/.github/workflows/cuda/Windows_env.sh
@@ -5,6 +5,7 @@ export CUDA_HOME="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v$CUDA_SHOR
 export PATH=$CUDA_HOME/bin:$PATH
 
 # https://github.com/pytorch/builder/blob/main/windows/cuda*.bat
+# https://github.com/pytorch/pytorch/blob/main/.ci/pytorch/windows/cuda*.bat
 declare -A CUDA_ARCHS=(
   ["cu102"]="3.7+PTX;5.0;6.0;6.1;7.0;7.5"
 
@@ -16,6 +17,7 @@ declare -A CUDA_ARCHS=(
 
   ["cu121"]="5.0;6.0;6.1;7.0;7.5;8.0;8.6;9.0"
   ["cu124"]="5.0;6.0;6.1;7.0;7.5;8.0;8.6;9.0"
+  ["cu126"]="5.0;6.0;6.1;7.0;7.5;8.0;8.6;9.0"
 )
 export TORCH_CUDA_ARCH_LIST=${CUDA_ARCHS[$CUDA_VERSION]}
 export FORCE_CUDA=1
diff --git a/.github/workflows/cuda/Windows_install.sh b/.github/workflows/cuda/Windows_install.sh
@@ -19,6 +19,7 @@ declare -A CUDA_LINKS=(
   ["cu123"]="https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_546.12_windows.exe"
   ["cu124"]="https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_551.78_windows.exe"
   ["cu125"]="https://developer.download.nvidia.com/compute/cuda/12.5.1/local_installers/cuda_12.5.1_555.85_windows.exe"
+  ["cu126"]="https://developer.download.nvidia.com/compute/cuda/12.6.0/local_installers/cuda_12.6.0_560.76_windows.exe"
 )
 CUDA_LINK=${CUDA_LINKS[$CUDA_VERSION]}
 

diff --git a/README.md b/README.md
@@ -60,7 +60,7 @@ pytorch3d==0.7.6+pt2.2.1cu121
 
 ## Supported combinations
 
-Tested with PyTorch `1.11.0` - `2.5.1` and their respective compute platforms and supported OSes, with an exception for `cu102` on Windows (no VS 2017 on the GH `windows-2019` runner), and the `rocm` platform.
+Tested with PyTorch `1.11.0` - `2.6.0` and their respective compute platforms and supported OSes, with an exception for `cu102` on Windows (no VS 2017 on the GH `windows-2019` runner), and the `rocm` platform.
 
 ## Pitfalls
 

diff --git a/get_jobs.py b/get_jobs.py
@@ -12,6 +12,7 @@ class TorchRelease:
 
 LINUX_X64 = "ubuntu-20.04"
 WINDOWS_X64 = "windows-2019"
+WINDOWS_X64_2022 = "windows-2022"
 MACOS_X64 = "macos-13"
 MACOS_ARM64 = "macos-14"
 
@@ -35,6 +36,7 @@ class TorchRelease:
     "2.4.1": TorchRelease(("3.8", "3.9", "3.10", "3.11", "3.12"), ("cpu", "cu118", "cu121", "cu124", "rocm6.1")),
     "2.5.0": TorchRelease(("3.9", "3.10", "3.11", "3.12", "3.13"), ("cpu", "cu118", "cu121", "cu124", "rocm6.2")),
     "2.5.1": TorchRelease(("3.9", "3.10", "3.11", "3.12", "3.13"), ("cpu", "cu118", "cu121", "cu124", "rocm6.2")),
+    "2.6.0": TorchRelease(("3.9", "3.10", "3.11", "3.12", "3.13"), ("cpu", "cu118", "cu124", "cu126", "rocm6.2.4")),
 }