is_gpu_available.py

from functools import lru_cache
import nvidia_smi


@lru_cache()
def is_gpu_available() -> bool:
    """Check if nvidia-smi is available."""
    try:
        nvidia_smi.nvmlInit()
        nvidia_smi.nvmlShutdown()
    except Exception:
        return False
    return True