import torch
import torch.version

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available:  {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"CUDA version:    {torch.version.cuda}")
    print(f"Device count:    {torch.cuda.device_count()}")
    print(f"Device name:     {torch.cuda.get_device_name(0)}")

    # Run a small tensor op on the GPU
    a = torch.randn(1000, 1000, device="cuda")
    b = torch.randn(1000, 1000, device="cuda")
    c = a @ b
    torch.cuda.synchronize()
    print(f"Matrix multiply:  OK  (result shape {c.shape}, device {c.device})")
else:
    print("No CUDA device found — check driver/CUDA installation.")