Windows AMI - Cuda 12.6 (#5880)

related to pytorch/pytorch#138440 cc @atalman
pytorch · Nov 9, 2024 · fbebf82 · fbebf82
1 parent 42c66e9
commit fbebf82
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/aws/ami/windows/scripts/Installers/Install-CUDA-Tools.ps1 b/aws/ami/windows/scripts/Installers/Install-CUDA-Tools.ps1
@@ -34,6 +34,11 @@ Switch ($cudaVersion) {
     $cudnn_subfolder = "cudnn-windows-x86_64-9.1.0.70_cuda12-archive"
     $installerArgs += " cuda_profiler_api_$cudaVersion nvjitlink_$cudaVersion"
   }
+  "12.6" {
+    $toolkitInstaller = "cuda_12.6.2_560.94_windows.exe"
+    $cudnn_subfolder = "cudnn-windows-x86_64-9.5.0.50_cuda12-archive"
+    $installerArgs += " cuda_profiler_api_$cudaVersion nvjitlink_$cudaVersion"
+  }
 }
 
 

diff --git a/aws/ami/windows/windows.pkr.hcl b/aws/ami/windows/windows.pkr.hcl
@@ -120,6 +120,13 @@ build {
     ]
   }
 
+  provisioner "powershell" {
+    environment_vars = ["CUDA_VERSION=12.6"]
+    scripts = [
+      "${path.root}/scripts/Installers/Install-CUDA-Tools.ps1",
+    ]
+  }
+
   # Uninstall Windows Defender, it brings more trouble than it's worth. Do this
   # last as it screws up the installation of other services like sshd somehow
   provisioner "powershell" {