check-vram.sh

#!/bin/bash
MB_SCALE=1024

GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -n 1)
GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
VRAM_MB=$(nvidia-smi -q -i 0 | grep Total | head -n 1 | awk '{ print $3 }')
VRAM_GB=$((VRAM_MB / MB_SCALE))

if [ -z "$VRAM_MB" ]; then
  >&2 echo "Graphics cards not detected by nvidia-smi!"
  exit 1
fi

>&2 echo "GPU $GPU_NAME * $GPU_COUNT (VRAM ${VRAM_GB}GiB) detected"
echo $VRAM_GB