RuntimeError: module must have its parameters and buffers on device cuda:2 (device_ids[0]) but found

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

出现这个问题是使用一机多卡分布式训练时，没有调用显卡0，例如使用显卡1,2,3进行训练。
使用的模块是pytorch nn.DataParallel。

下面给出有问题的代码

    device = torch.device("cuda" if args.cuda else "cpu")
    #Try to load models
    if args.model