From a77b4008f6eafc8ec6d83bf7bd128575b12dfabb Mon Sep 17 00:00:00 2001
From: Steel <79706171+Steel-skull@users.noreply.github.com>
Date: Wed, 13 Nov 2024 22:10:26 -0600
Subject: [PATCH] Update gpu-switch.sh

---
 gpu-switch.sh | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/gpu-switch.sh b/gpu-switch.sh
index 31f686c..22238e8 100644
--- a/gpu-switch.sh
+++ b/gpu-switch.sh
@@ -35,21 +35,24 @@ convert_to_pci_address() {
     local device="$1"
     local gpu_address=""
 
-    if [[ "$device" =~ ^[0-9]+$ ]]; then
-        # Convert GPU index to PCI address
-        gpu_address=$(nvidia-smi --id=$device --query-gpu=gpu_bus_id --format=csv,noheader 2>/dev/null | tr -d '[:space:]')
-    elif [[ "$device" =~ ^GPU-.*$ ]]; then
-        # Handle UUID
-        gpu_address=$(nvidia-smi --id=$device --query-gpu=gpu_bus_id --format=csv,noheader 2>/dev/null | tr -d '[:space:]')
+    if [[ "$device" =~ ^[0-9]+$ || "$device" =~ ^GPU-.*$ ]]; then
+        # Convert GPU index or UUID to PCI address
+        gpu_address=$(nvidia-smi --id="$device" --query-gpu=gpu_bus_id --format=csv,noheader 2>/dev/null | tr -d '[:space:]')
     else
         # Direct PCI address provided
-        gpu_address=$device
+        gpu_address="$device"
+    fi
+
+    # Check for valid output
+    if [ -z "$gpu_address" ]; then
+        error_exit "Failed to get PCI address for device: $device"
     fi
 
     # Standardize format
-    echo "$gpu_address" | sed 's/0000://' | sed 's/\./:/g'
+    echo "$gpu_address" | sed -e 's/0000://' -e 's/\./:/g'
 }
 
+
 get_gpu_addresses() {
     # Split devices by comma
     IFS=',' read -ra DEVICES <<< "$NVIDIA_VISIBLE_DEVICES"