scripts/hf_jobs/gpu_test.py

# /// script
# requires-python = ">=3.12"
# dependencies = ["torch>=2.5.0"]
# ///
import torch, sys
print(f"Python: {sys.version}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"Mem: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
print("SUCCESS")