SpiceNet / code /_gpu_check.py
Noushad999's picture
Upload folder using huggingface_hub
1ea7ba6 verified
Raw
History Blame Contribute Delete
458 Bytes
import torch
print(f"torch={torch.__version__}")
print(f"cuda={torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"dev={torch.cuda.get_device_name(0)}")
cap = torch.cuda.get_device_capability(0)
print(f"cap=sm_{cap[0]}{cap[1]}")
print(f"mem={torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB")
x = torch.randn(2048, 2048).cuda()
y = (x @ x).cpu()
print(f"matmul OK mean={y.mean().item():.4f}")