pytorch调用GPU的代码

import torch

检查是否有可用的GPU

if torch.cuda.is_available():
# 创建一个设备对象，使用默认的GPU设备
device = torch.device(“cuda”)
else:
# 如果没有可用的GPU，则使用CPU设备
device = torch.device(“cpu”)

创建一个张量，并将其移动到设备上

x = torch.tensor([1, 2, 3])
x = x.to(device)

创建一个模型，并将其移动到设备上

model = YourModel()
model = model.to(device)

在设备上进行前向传播

output = model(x)

将输出移回到CPU，并将其转换为NumPy数组（如果需要）

output = output.to(“cpu”).detach().numpy()