两张显卡如何并行进行cuda加速，请给出pytorch示例

在 PyTorch 中，可以通过设置 `CUDA_VISIBLE_DEVICES` 环境变量来选择使用哪些显卡，然后使用 `torch.nn.DataParallel` 模块将模型并行化。以下是一个示例：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader

# 设置使用的显卡
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
    torch.cuda.set_device([0, 1])  # 选择第 0 和 1 号显卡进行计算
print("Using device:", device)

# 定义模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = Model().to(device)

# 并行化模型
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

# 定义数据集和 DataLoader
train_dataset = torch.randn(1000, 10)
train_labels = torch.randn(1000, 1)
train_loader = DataLoader(list(zip(train_dataset, train_labels)), batch_size=32)

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(10):
    for i, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        print("Epoch:", epoch, "Batch:", i, "Loss:", loss.item())

在这个示例中，我们首先设置了环境变量 `CUDA_VISIBLE_DEVICES` 来选择使用第 0 和 1 号显卡进行计算。然后，我们定义了一个简单的神经网络模型 `Model`，它包含两个线性层。接下来，我们使用 `nn.DataParallel` 将模型并行化，如果只有一张显卡则不需要并行化。我们还定义了一个简单的数据集和 DataLoader，并使用均方误差损失函数和随机梯度下降优化器来训练模型。在训练过程中，我们将输入和标签移动到选择的显卡上进行计算。