Add back dict

pytorch · Feb 3, 2025 · 862bf9d · 862bf9d
1 parent 4c0208c
commit 862bf9d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py
@@ -37,6 +37,13 @@
     "release": ["6.1", "6.2.4"],
 }
 
+CUDA_CUDNN_VERSIONS = {
+    "11.8": {"cuda": "11.8.0", "cudnn": "9"},
+    "12.4": {"cuda": "12.4.1", "cudnn": "9"},
+    "12.6": {"cuda": "12.6.3", "cudnn": "9"},
+    "12.8": {"cuda": "12.8.0", "cudnn": "9"},
+}
+
 PACKAGE_TYPES = ["wheel", "conda", "libtorch"]
 PRE_CXX11_ABI = "pre-cxx11"
 CXX11_ABI = "cxx11-abi"

diff --git a/tools/scripts/generate_docker_release_matrix.py b/tools/scripts/generate_docker_release_matrix.py
@@ -45,7 +45,7 @@ def generate_docker_matrix(
         docker_image_version = f"{prefix}-nightly:{generate_binary_build_matrix.CURRENT_NIGHTLY_VERSION}.dev{datetime.today().strftime('%Y%m%d')}"
 
     for cuda in generate_binary_build_matrix.CUDA_ARCHES_DICT[channel]:
-        version = generate_binary_build_matrix.CUDA_CUDDN_VERSIONS[cuda]
+        version = generate_binary_build_matrix.CUDA_CUDNN_VERSIONS[cuda]
         for image in DOCKER_IMAGE_TYPES:
             ret.append(
                 {