Fix building wheels for torch 2.2.0. (#95)

See also pytorch/pytorch#120020
csukuangfj · Feb 18, 2024 · 2a8a993 · 2a8a993
1 parent 2e042b3
commit 2a8a993
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/.github/workflows/ubuntu-cpu-wheels.yml b/.github/workflows/ubuntu-cpu-wheels.yml
@@ -28,6 +28,9 @@ jobs:
           # outputting for debugging purposes
           python ./scripts/github_actions/generate_build_matrix.py
           MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py)
+
+          # python ./scripts/github_actions/generate_build_matrix.py --test-only-latest-torch
+          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --test-only-latest-torch)
           echo "::set-output name=matrix::${MATRIX}"
 
   build-manylinux-wheels:

diff --git a/scripts/github_actions/generate_build_matrix.py b/scripts/github_actions/generate_build_matrix.py
@@ -5,6 +5,18 @@
 import json
 
 
+def version_ge(a, b):
+    a_major, a_minor = list(map(int, a.split(".")))[:2]
+    b_major, b_minor = list(map(int, b.split(".")))[:2]
+    if a_major > b_major:
+        return True
+
+    if a_major == b_major and a_minor >= b_minor:
+        return True
+
+    return False
+
+
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -216,6 +228,14 @@ def generate_build_matrix(enable_cuda, for_windows, for_macos, test_only_latest_
                     ans.append({"torch": torch, "python-version": p})
                 elif for_macos:
                     ans.append({"torch": torch, "python-version": p})
+                elif version_ge(torch, "2.2.0"):
+                    ans.append(
+                        {
+                            "torch": torch,
+                            "python-version": p,
+                            "image": "pytorch/manylinux-builder:cpu-2.2",
+                        }
+                    )
                 else:
                     ans.append(
                         {