19 lines
654 B
Python
19 lines
654 B
Python
import torch
|
|
|
|
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.")
|