torch指定GPU训练

要在指定的GPU上训练PyTorch模型，可以使用以下代码：

import torch

# 指定要使用的GPU编号
device = torch.device("cuda:0")

# 将模型和数据移动到指定的GPU
model.to(device)
data = data.to(device)

# 在指定的GPU上进行训练
for epoch in range(num_epochs):
    for batch in data_loader:
        optimizer.zero_grad()
        inputs, targets = batch
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        loss.backward()
        optimizer.step()

在这个例子中，我们将模型和数据移动到GPU `cuda:0` 上，然后在这个GPU上进行训练。要使用其他GPU，只需要将 `cuda:0` 替换为相应的GPU编号即可。