Fabian
30 lines
911 B
Python
30 lines
911 B
Python
import torch
|
|
import os
|
|
|
|
print(f"HSA_OVERRIDE_GFX_VERSION={os.environ.get('HSA_OVERRIDE_GFX_VERSION','NOT SET')}")
|
|
print(f"PyTorch arch list: {torch.cuda.get_arch_list()}")
|
|
print(f"Device: {torch.cuda.get_device_name(0)}")
|
|
print(f"GCN Arch: {torch.cuda.get_device_properties(0).gcnArchName}")
|
|
|
|
try:
|
|
t = torch.randn(256, 256, device="cuda")
|
|
r = torch.mm(t, t)
|
|
val = r[0,0].item()
|
|
print(f"GPU matmul OK! result[0,0]={val:.4f}")
|
|
|
|
# Bigger test
|
|
a = torch.randn(1024, 1024, device="cuda")
|
|
b = torch.randn(1024, 1024, device="cuda")
|
|
c = torch.mm(a, b)
|
|
print(f"Large matmul OK! shape={c.shape}")
|
|
|
|
# Test fp32 conv
|
|
x = torch.randn(1, 3, 64, 64, device="cuda")
|
|
conv = torch.nn.Conv2d(3, 16, 3, padding=1).cuda()
|
|
y = conv(x)
|
|
print(f"Conv2d OK! output shape={y.shape}")
|
|
|
|
print("ALL GPU TESTS PASSED!")
|
|
except Exception as e:
|
|
print(f"FAILED: {e}")
|