diff --git a/misc/tools.func b/misc/tools.func
index 6e407ddce..447246a73 100644
--- a/misc/tools.func
+++ b/misc/tools.func
@@ -4685,10 +4685,10 @@ _setup_nvidia_gpu() {
   # Format varies by driver type:
   #   Proprietary: "NVRM version: NVIDIA UNIX x86_64 Kernel Module  550.54.14  Thu..."
   #   Open:        "NVRM version: NVIDIA UNIX Open Kernel Module for x86_64  590.48.01  Release..."
-  # Use regex to extract version number (###.##.## pattern)
+  # Use regex to extract version number (###.##.## or ###.## pattern)
   local nvidia_host_version=""
   if [[ -f /proc/driver/nvidia/version ]]; then
-    nvidia_host_version=$(grep -oP '\d{3,}\.\d+\.\d+' /proc/driver/nvidia/version 2>/dev/null | head -1)
+    nvidia_host_version=$(grep -oP '\d{3,}\.\d+(\.\d+)?' /proc/driver/nvidia/version 2>/dev/null | head -1)
   fi
 
   if [[ -z "$nvidia_host_version" ]]; then