| 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}") | |