File size: 798 Bytes
0f96bb5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | import torch
print('PyTorch version:', torch.__version__)
print('CUDA available:', torch.cuda.is_available())
if torch.cuda.is_available():
print('CUDA device:', torch.cuda.get_device_name(0))
print('CUDA version:', torch.version.cuda)
print('VRAM:', round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2), 'GB')
# Test actual GPU computation
try:
x = torch.randn(100, 100, device='cuda')
y = torch.randn(100, 100, device='cuda')
z = torch.mm(x, y)
print('GPU computation test: PASSED')
print('Result shape:', z.shape, 'sum:', z.sum().item())
except Exception as e:
print('GPU computation test: FAILED -', e)
print('Falling back to CPU mode')
else:
print('WARNING: CUDA not available, will use CPU')
|