From 32e40d3a7fa90fded5d468e6a07f1aa4f2fe3ef4 Mon Sep 17 00:00:00 2001
From: Evan Hosinski <ehosinski@krkn.tech>
Date: Thu, 22 Jan 2026 11:54:36 -0500
Subject: [PATCH] Fixed the Cuda detection and fallback logic

On systems where CUDA_VERSION=12.4 the script was incorrectly labeling it as 12.6 by default
---
 install.sh | 29 ++++++++++++++++-------------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/install.sh b/install.sh
index 054743f..68e5dd1 100644
--- a/install.sh
+++ b/install.sh
@@ -88,21 +88,24 @@ fi
 DETECTED_CUDA="$(nvidia-smi | grep 'CUDA Version' | sed 's/.*CUDA Version: *\([0-9.]*\).*/\1/')"
 log "Detected CUDA capability: ${DETECTED_CUDA}"
 
-choose_cuda() {
-  for v in "${SUPPORTED_CUDA[@]}"; do
-    if [[ "$DETECTED_CUDA" == "$v"* ]]; then
-      echo "$v"
-      return
-    fi
-  done
+  choose_cuda() {
+    # Pick highest supported version <= detected CUDA capability
+    for v in $(printf '%s\n' "${SUPPORTED_CUDA[@]}" | sort -rV); do
+      if [ "$(printf '%s\n%s\n' "$v" "$DETECTED_CUDA" | sort -V | head -n1)" = "$v" ]; then
+        echo "$v"
+        return
+      fi
+    done
+    echo ""
+  }
 
   # fallback: highest <= detected
-  for v in $(printf '%s\n' "${SUPPORTED_CUDA[@]}" | sort -rV); do
-    if [[ "$(printf '%s\n%s\n' "$v" "$DETECTED_CUDA" | sort -V | head -n1)" == "$v" ]]; then
-      echo "$v"
-      return
-    fi
-  done
+  if [ -z "$CUDA_VERSION" ]; then
+    warn "Could not auto-match CUDA version, defaulting to highest supported"
+    CUDA_VERSION="$(printf '%s\n' "${SUPPORTED_CUDA[@]}" | sort -V | tail -n1)"
+  else
+    log "Selected CUDA image version: ${CUDA_VERSION}"
+  fi
 
   echo ""
 }