From 7bce4cb93b6335bf24d054cb86e3263be6c04c35 Mon Sep 17 00:00:00 2001
From: Matthew Douglas <38992547+matthewdouglas@users.noreply.github.com>
Date: Fri, 29 Nov 2024 10:55:55 -0500
Subject: [PATCH] [Build] Add CUDA 12.6.2 build; update 12.5.0 to 12.5.1

---
 .github/workflows/python-package.yml | 2 +-
 install_cuda.py                      | 3 ++-
 install_cuda.sh                      | 6 +++++-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index 560741edb..d9ff1dda8 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -64,7 +64,7 @@ jobs:
         os: [ubuntu-latest, windows-latest]
         arch: [x86_64, aarch64]
         cuda_version:
-          ["11.7.1", "11.8.0", "12.0.1", "12.1.1", "12.2.2", "12.3.2", "12.4.1", "12.5.0"]
+          ["11.7.1", "11.8.0", "12.0.1", "12.1.1", "12.2.2", "12.3.2", "12.4.1", "12.5.1", "12.6.2"]
         exclude:
           - os: windows-latest # This probably requires arm64 Windows agents
             arch: aarch64
diff --git a/install_cuda.py b/install_cuda.py
index 8267c5e2b..a1fe2c545 100644
--- a/install_cuda.py
+++ b/install_cuda.py
@@ -18,7 +18,8 @@
     "122": "https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run",
     "123": "https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.run",
     "124": "https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run",
-    "125": "https://developer.download.nvidia.com/compute/cuda/12.5.0/local_installers/cuda_12.5.0_555.42.02_linux.run",
+    "125": "https://developer.download.nvidia.com/compute/cuda/12.5.1/local_installers/cuda_12.5.1_555.42.06_linux.run",
+    "126": "https://developer.download.nvidia.com/compute/cuda/12.6.2/local_installers/cuda_12.6.2_560.35.03_linux.run",
 }
 
 
diff --git a/install_cuda.sh b/install_cuda.sh
index 0aa9531fc..269981f66 100644
--- a/install_cuda.sh
+++ b/install_cuda.sh
@@ -12,7 +12,8 @@ URL121=https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installer
 URL122=https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run
 URL123=https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.run
 URL124=https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run
-URL125=https://developer.download.nvidia.com/compute/cuda/12.5.0/local_installers/cuda_12.5.0_555.42.02_linux.run
+URL125=https://developer.download.nvidia.com/compute/cuda/12.5.1/local_installers/cuda_12.5.1_555.42.06_linux.run
+URL126=https://developer.download.nvidia.com/compute/cuda/12.6.2/local_installers/cuda_12.6.2_560.35.03_linux.run
 
 CUDA_VERSION=$1
 BASE_PATH=$2
@@ -64,6 +65,9 @@ if [[ -n "$CUDA_VERSION" ]]; then
   elif [[ "$CUDA_VERSION" -eq "125" ]]; then
     URL=$URL125
     FOLDER=cuda-12.5
+  elif [[ "$CUDA_VERSION" -eq "126" ]]; then
+    URL=$URL126
+    FOLDER=cuda-12.6
   else
     echo "argument error: No cuda version passed as input. Choose among versions 110 to 125"
   fi