Fabian
8 lines
172 B
Python
8 lines
172 B
Python
import torch
|
|
print("CUDA:", torch.cuda.is_available())
|
|
t = torch.randn(4, 4, device="cuda")
|
|
print("GPU tensor:", t.shape)
|
|
r = t @ t.T
|
|
print("matmul:", r.shape)
|
|
print("OK")
|