From d0091a7bfe576bcd74caae5f584067a4e5889525 Mon Sep 17 00:00:00 2001
From: Bryce Soghigian <49734722+Bryce-Soghigian@users.noreply.github.com>
Date: Mon, 8 Jan 2024 18:20:14 -0800
Subject: [PATCH] chore: bump cuda gpu driver version from 525 to 535 (#85)

---
 pkg/utils/gpu.go      | 6 +++---
 pkg/utils/gpu_test.go | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/pkg/utils/gpu.go b/pkg/utils/gpu.go
index e66fbdbfb..a6b657139 100644
--- a/pkg/utils/gpu.go
+++ b/pkg/utils/gpu.go
@@ -23,11 +23,11 @@ import (
 // TODO: Get these from agentbaker
 const (
 	Nvidia470CudaDriverVersion = "cuda-470.82.01"
-	Nvidia525CudaDriverVersion = "cuda-525.85.12"
+	Nvidia535CudaDriverVersion = "cuda-535.54.03"
 	Nvidia535GridDriverVersion = "grid-535.54.03"
 
 	AKSGPUGridSHA = "sha-20ffa2"
-	AKSGPUCudaSHA = "sha-e8873b"
+	AKSGPUCudaSHA = "sha-ff213d"
 )
 
 func GetAKSGPUImageSHA(size string) string {
@@ -146,7 +146,7 @@ func GetGPUDriverVersion(size string) string {
 	if isStandardNCv1(size) {
 		return Nvidia470CudaDriverVersion
 	}
-	return Nvidia525CudaDriverVersion
+	return Nvidia535CudaDriverVersion
 }
 
 func isStandardNCv1(size string) bool {
diff --git a/pkg/utils/gpu_test.go b/pkg/utils/gpu_test.go
index 71e683fa7..40a70db21 100644
--- a/pkg/utils/gpu_test.go
+++ b/pkg/utils/gpu_test.go
@@ -55,9 +55,9 @@ func TestGetGPUDriverVersion(t *testing.T) {
 	}{
 		{"GRID Driver - NV Series v5", "standard_nv6ads_a10_v5", Nvidia535GridDriverVersion},
 		{"CUDA Driver - NC Series v1", "standard_nc6s", Nvidia470CudaDriverVersion},
-		{"CUDA Driver - NC Series v2", "standard_nc6s_v2", Nvidia525CudaDriverVersion},
-		{"Unknown SKU", "unknown_sku", Nvidia525CudaDriverVersion},
-		{"CUDA Driver - NC Series v3", "standard_nc6s_v3", Nvidia525CudaDriverVersion},
+		{"CUDA Driver - NC Series v2", "standard_nc6s_v2", Nvidia535CudaDriverVersion},
+		{"Unknown SKU", "unknown_sku", Nvidia535CudaDriverVersion},
+		{"CUDA Driver - NC Series v3", "standard_nc6s_v3", Nvidia535CudaDriverVersion},
 		{"GRID Driver - A10", "standard_nc8ads_a10_v4", Nvidia535GridDriverVersion},
 	}