File size: 458 Bytes
1ea7ba6 | 1 2 3 4 5 6 7 8 9 10 11 12 | 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}")
|