module must have its parameters and buffers on device cuda:1 (device_ids[0]) but found one of them o

一机多卡训练时报错：RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cuda:0
假设我们设置在1，2，3，4这个四张gpu上训练
原来代码为：

device_ids = [1,2,3,4]
model = torch.nn.DataParallel(model, device_ids=device_ids)

将其改成：

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3,4"
device_ids = [0,1,2,3]
model = torch.nn.DataParallel(model, device_ids=device_ids)

即可解决