Fix adding cuda 12.6 logic

pytorch · Nov 18, 2024 · cf8a3d0 · cf8a3d0
1 parent 661fe3c
commit cf8a3d0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py
@@ -517,9 +517,9 @@ def generate_wheels_matrix(
             upload_to_base_bucket = "no"
             if os == LINUX or os == WINDOWS:
                 arches += CUDA_ARCHES
-            # todo: remove once windows cuda 12.6 binaries are available
-            if channel == NIGHTLY and os != LINUX:
-                arches.remove("12.6")
+                # todo: remove once windows cuda 12.6 binaries are available
+                if channel == NIGHTLY and os == WINDOWS:
+                    arches.remove("12.6")
 
         if with_rocm == ENABLE:
             if os == LINUX: