Add a dev image for cuda12.3 (pytorch#6835)

lsy323 · Mar 27, 2024 · 59db0cd · 59db0cd
1 parent d99ceb6
commit 59db0cd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/infra/ansible/config/cuda_deps.yaml b/infra/ansible/config/cuda_deps.yaml
@@ -3,12 +3,14 @@
 cuda_deps:
   # List all libcudnn8 versions with `apt list -a libcudnn8`
   libcudnn:
+    "12.3": libcudnn9-cuda-12=9.0.0.312-1
     "12.1": libcudnn8=8.9.2.26-1+cuda12.1
     "12.0": libcudnn8=8.8.0.121-1+cuda12.0
     "11.8": libcudnn8=8.7.0.84-1+cuda11.8
     "11.7": libcudnn8=8.5.0.96-1+cuda11.7
     "11.2": libcudnn8=8.1.1.33-1+cuda11.2
   libcudnn-dev:
+    "12.3": libcudnn9-dev-cuda-12=9.0.0.312-1
     "12.1": libcudnn8-dev=8.9.2.26-1+cuda12.1
     "12.0": libcudnn8-dev=8.8.0.121-1+cuda12.0
     "11.8": libcudnn8-dev=8.7.0.84-1+cuda11.8

diff --git a/infra/tpu-pytorch-releases/dev_images.auto.tfvars b/infra/tpu-pytorch-releases/dev_images.auto.tfvars
@@ -13,5 +13,11 @@ dev_images = [
     accelerator  = "cuda"
     cuda_version = "12.1"
     extra_tags   = ["cuda"]
+  },
+  {
+    accelerator  = "cuda"
+    cuda_version = "12.3"
+    extra_tags   = ["cuda"]
+    python_version = "3.10"
   }
 ]