我有多个 GPU 设备,想在它们上运行 Pytorch。我已经在我的代码中尝试了多 GPU 示例和数据并行
device = torch.device("cuda:0,1,2")
model = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
model.to(device)
Run Code Online (Sandbox Code Playgroud)
在我的代码中。但训练仍然在一个 GPU 上进行(cuda:0)。在 shell 中,我还在export CUDA_VISIBLE_DEVICES=0,1,2运行代码之前选择了使用的 GPU。
有人可以帮我解决这个问题吗?