pytorch gpu测试代码

以下是一个使用 PyTorch 在 GPU 上运行的简单测试代码：

import torch

# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using {device} device")

# 创建随机张量并将其移动到设备上
x = torch.rand(3, 3).to(device)

# 进行张量操作
y = torch.rand(3, 3).to(device)
z = torch.matmul(x, y)

# 将结果移回 CPU 并打印
result = z.to("cpu")
print(result)

在此代码中，我们首先检查是否可用 GPU。然后，我们创建了一个随机张量并将其移动到设备上。接下来，我们执行一些张量操作，将结果移回 CPU 并打印。