Use cuda 12.6 wheels with Manylinux 2.28. Use Manylinux2014 for CPU, …

…CUDA11.8, CUDA12.4 (#5986) PyTorch PR: https://github.com/pytorch/pytorch/pull/141565/files
pytorch · Nov 26, 2024 · 8856c07 · 8856c07
1 parent f75268c
commit 8856c07
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py
@@ -148,15 +148,14 @@ def initialize_globals(channel: str, build_python_only: bool) -> None:
     else:
         PYTHON_ARCHES = PYTHON_ARCHES_DICT[channel]
     WHEEL_CONTAINER_IMAGES = {
-        **{
-            gpu_arch: f"pytorch/manylinux2_28-builder:cuda{gpu_arch}"
-            for gpu_arch in CUDA_ARCHES
-        },
+       "11.8": "pytorch/manylinux-builder:cuda11.8",
+       "12.4": "pytorch/manylinux-builder:cuda12.4",
+       "12.6": "pytorch/manylinux2_28-builder:cuda12.6",
         **{
             gpu_arch: f"pytorch/manylinux-builder:rocm{gpu_arch}"
             for gpu_arch in ROCM_ARCHES
         },
-        CPU: "pytorch/manylinux2_28-builder:cpu",
+        CPU: "pytorch/manylinux-builder:cpu",
         XPU: "pytorch/manylinux2_28-builder:xpu",
         # TODO: Migrate CUDA_AARCH64 image to manylinux2_28_aarch64-builder:cuda12.4
         CPU_AARCH64: "pytorch/manylinux2_28_aarch64-builder:cpu-aarch64",