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")