如何在多个 GPU 上运行 Pytorch 代码？

Question

我有多个 GPU 设备，想在它们上运行 Pytorch。我已经在我的代码中尝试了多 GPU 示例和数据并行

device = torch.device("cuda:0,1,2")
model = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
model.to(device)

在我的代码中。但训练仍然在一个 GPU 上进行（cuda:0）。在 shell 中，我还在export CUDA_VISIBLE_DEVICES=0,1,2运行代码之前选择了使用的 GPU。

有人可以帮我解决这个问题吗？