举例说明 torch.cuda.set_device() 如何指定多张GPU

torch.cuda.set_device() 函数用于设置当前使用的 GPU，可以指定一张或多张 GPU。下面是一个指定多张 GPU 的例子：

import torch

# 指定使用的 GPU
device_ids = [0, 1, 2, 3]  # 假设有四张 GPU
torch.cuda.set_device(device_ids[0])  # 设置当前使用的第一张 GPU

# 定义模型并将模型放到多张 GPU 上
model = torch.nn.DataParallel(model, device_ids=device_ids)
model.cuda()

在上面的例子中，我们通过 `device_ids` 列表指定了要使用的多张 GPU，然后将第一张 GPU 设置为当前使用的 GPU。接着，我们将模型使用 `torch.nn.DataParallel` 并指定 `device_ids` 参数来实现在多张 GPU 上并行计算。最后，我们将模型移动到 GPU 上。这样，我们就可以在多张 GPU 上并行计算模型了。