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