From 1cf8844dec07f5644742b54aa37cf812a46ca965 Mon Sep 17 00:00:00 2001
From: Fangjun Kuang <csukuangfj@gmail.com>
Date: Thu, 30 Jan 2025 08:48:53 +0800
Subject: [PATCH] Support torch 2.6.0 (#1324)

---
 .github/workflows/windows-x64-cpu-wheels.yml    |  2 +-
 scripts/github_actions/generate_build_matrix.py | 12 ++++++++++--
 scripts/github_actions/install_torch.sh         | 17 +++++++++++++++++
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/windows-x64-cpu-wheels.yml b/.github/workflows/windows-x64-cpu-wheels.yml
index 7bda39b62..1884a5ee6 100644
--- a/.github/workflows/windows-x64-cpu-wheels.yml
+++ b/.github/workflows/windows-x64-cpu-wheels.yml
@@ -49,7 +49,7 @@ jobs:
       # see https://cibuildwheel.readthedocs.io/en/stable/changelog/
       # for a list of versions
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.16.5
+        uses: pypa/cibuildwheel@v2.22.0
         env:
           CIBW_BEFORE_BUILD: pip install torch==${{ matrix.torch}}+cpu cmake numpy -f https://download.pytorch.org/whl/torch_stable.html || pip install torch==${{ matrix.torch}}+cpu cmake numpy -f https://download.pytorch.org/whl/torch/
           CIBW_BUILD: ${{ matrix.python-version }}-win_amd64
diff --git a/scripts/github_actions/generate_build_matrix.py b/scripts/github_actions/generate_build_matrix.py
index abe1582ad..d455c9fc4 100755
--- a/scripts/github_actions/generate_build_matrix.py
+++ b/scripts/github_actions/generate_build_matrix.py
@@ -290,16 +290,24 @@ def generate_build_matrix(
                 else ["11.8.0", "12.1.0", "12.4.0"]
             ),
         },
+        "2.6.0": {
+            "python-version": ["3.9", "3.10", "3.11", "3.12", "3.13"],
+            "cuda": (
+                ["11.8", "12.4", "12.6"]  # default 12.4
+                if not for_windows
+                else ["11.8.0", "12.4.0", "12.6.0"]
+            ),
+        },
         # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/windows-links.ts
     }
     if test_only_latest_torch:
-        latest = "2.5.1"
+        latest = "2.6.0"
         matrix = {latest: matrix[latest]}
 
     if for_windows or for_macos:
         if "2.5.1" in matrix:
             matrix["2.5.1"]["python-version"].remove("3.13")
-            
+
         if "2.5.0" in matrix:
             matrix["2.5.0"]["python-version"].remove("3.13")
 
diff --git a/scripts/github_actions/install_torch.sh b/scripts/github_actions/install_torch.sh
index 4d06b90cc..80de75ad1 100755
--- a/scripts/github_actions/install_torch.sh
+++ b/scripts/github_actions/install_torch.sh
@@ -247,6 +247,23 @@ case ${torch} in
         ;;
     esac
     ;;
+  2.6.*)
+    case ${cuda} in
+      11.8)
+        package="torch==${torch}+cu118"
+        # https://download.pytorch.org/whl/nightly/torch/
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+      12.4)
+        package="torch==${torch}+cu124"
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+      12.6)
+        package="torch==${torch}+cu126"
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+    esac
+    ;;
   *)
     echo "Unsupported PyTorch version: ${torch}"
     exit 1