diff --git a/aio/entrypoint.sh b/aio/entrypoint.sh
index aeb5e4de..795cb86a 100755
--- a/aio/entrypoint.sh
+++ b/aio/entrypoint.sh
@@ -57,29 +57,33 @@ function detect_gpu() {
 }
 
 function detect_gpu_size() {
-    if [ "$GPU_ACCELERATION" = true ]; then
-        GPU_SIZE=gpu-8g
-    fi
-
     # Attempting to find GPU memory size for NVIDIA GPUs
-    if echo "$gpu_model" | grep -iq nvidia; then
+    if [ "$GPU_ACCELERATION" = true ] && [ "$GPU_VENDOR" = "nvidia" ]; then
         echo "NVIDIA GPU detected. Attempting to find memory size..."
-        nvidia_sm=($(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits))
+        # Using head -n 1 to get the total memory of the 1st NVIDIA GPU detected.
+        # If handling multiple GPUs is required in the future, this is the place to do it
+        nvidia_sm=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -n 1)
         if [ ! -z "$nvidia_sm" ]; then
-            echo "Total GPU Memory: ${nvidia_sm[0]} MiB"
+            echo "Total GPU Memory: $nvidia_sm MiB"
+            # if bigger than 8GB, use 16GB
+            #if [ "$nvidia_sm" -gt 8192 ]; then
+            #    GPU_SIZE=gpu-16g
+            #else
+            GPU_SIZE=gpu-8g
+            #fi
         else
-            echo "Unable to determine NVIDIA GPU memory size."
+            echo "Unable to determine NVIDIA GPU memory size. Falling back to CPU."
+            GPU_SIZE=gpu-8g
         fi
-        # if bigger than 8GB, use 16GB
-        #if [ "$nvidia_sm" -gt 8192 ]; then
-        #    GPU_SIZE=gpu-16g
-        #fi
-    else
-        echo "Non-NVIDIA GPU detected. GPU memory size detection for non-NVIDIA GPUs is not supported in this script."
-    fi
+
+    # Default to a generic GPU size until we implement GPU size detection for non NVIDIA GPUs
+    elif [ "$GPU_ACCELERATION" = true ]; then
+        echo "Non-NVIDIA GPU detected. Specific GPU memory size detection is not implemented."
+        GPU_SIZE=gpu-8g
 
     # default to cpu if GPU_SIZE is not set
-    if [ -z "$GPU_SIZE" ]; then
+    else
+        echo "GPU acceleration is not enabled or supported. Defaulting to CPU."
         GPU_SIZE=cpu
     fi
 }