env_setup/gpu_sanity_check.py

import torch

print(f"PyTorch version : {torch.__version__}")
print(f"CUDA available  : {torch.cuda.is_available()}")
print(f"GPU             : {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None — using CPU'}")

# Quick tensor op on GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.randn(3, 224, 224).to(device)
print(f"Test tensor on  : {x.device} — shape {x.shape}")