14 lines
477 B
Python
14 lines
477 B
Python
import torch
|
|
print(f"PyTorch: {torch.__version__}")
|
|
print(f"CUDA available: {torch.cuda.is_available()}")
|
|
print(f"HIP version: {torch.version.hip}")
|
|
print(f"Device count: {torch.cuda.device_count()}")
|
|
if torch.cuda.is_available():
|
|
print(f"Device name: {torch.cuda.get_device_name(0)}")
|
|
t = torch.randn(100, 100, device="cuda")
|
|
r = torch.mm(t, t)
|
|
print(f"GPU matmul OK: result shape {r.shape}")
|
|
print("GPU COMPUTE WORKS!")
|
|
else:
|
|
print("NO GPU DETECTED")
|