From cf8a3d0ee15d0579243b01679cccae272fd457eb Mon Sep 17 00:00:00 2001
From: atalman <atalman@fb.com>
Date: Mon, 18 Nov 2024 12:41:17 -0800
Subject: [PATCH] Fix adding cuda 12.6 logic

---
 tools/scripts/generate_binary_build_matrix.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py
index ab7cc214c7..950a9f446b 100644
--- 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: