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

Question

我有多个 GPU 设备并想在它们上运行 Pytorch。我已经在我的代码中尝试了 MULTI-GPU EXAMPLES 和 DATA PARALLELISM

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。

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