pytorch-gpu测试代码

我可以回答这个问题。您可以使用以下代码来测试您的 PyTorch GPU 是否正常工作：

import torch

if torch.cuda.is_available():
device = torch.device("cuda")
print("GPU is available")
else:
device = torch.device("cpu")
print("GPU is not available")

x = torch.randn(3, 3).to(device)
y = torch.randn(3, 3).to(device)
z = x + y

print(z)