pytorch指定GPU后仍旧使用GPU0的解决方法

我在训练模型的时候，已经指定了使用GPU5、6进行训练，然而pytorch却强行在GPU0里占用100多M的空间，但是刚好GPU0被别人占满了，导致一直是CUDA:out of memory的状态.
这个问题跟为什么Pytorch死活要用第0块显卡–我和pytorch的恩怨情仇非常像，但是我用他的方法根本不管用，感觉我的情况跟他的应该是不相同的，所以我只能自己找问题。

我根据错误提示，发现溢出的语句是加载预训练模型参数的地方：

rnet_checkpoint = torch.load('../s_resnet50_0.25_0.5_0.75_1.0.pt')['model']
rnet.load_state_dict(rnet_checkpoint)

原本我想的是，我前面已经把rnet放到了显卡5,6上：

rnet.to(device)
rnet = torch.nn.DataParallel(rnet, device_ids=args.deviceid)  # for multi gpu

那么我加载参数的时候也是加载到指定的GPU上才对，但是现在看来，是还没有来得及把参数加载到显卡上，pytorch可能是从checkpoint文件中提取出参数时，是默认把rnet_checkpoint 这个临时变量暂存在GPU0上的，后面才把rnet_checkpoint 加载进显卡5,6上的rnet中。
然而由于GPU0刚好没位置导致出错，所以我后面加了一个参数强制把rnet_checkpoint 放到cpu上暂存，这样就不会默认放到GPU0上了

rnet_checkpoint = torch.load('../s_resnet50_0.25_0.5_0.75_1.0.pt',map_location='cpu')['model']
rnet.load_state_dict(rnet_checkpoint)