如何在pytorch框架中使用多gpu进行推理

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(10, 1)
    
    def forward(self, x):
        return self.fc(x)

model = MyModel()
model = nn.DataParallel(model)  # 将模型包装为DataParallel模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

input_data = torch.randn(16, 10).to(device)

output = model(input_data)