pytorch并行计算，显存溢出问题之torch并行问题解决办法

使用的是torch模型，遇到一个显存超出报错问题，信息如下：

  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 443, in to
    return self._apply(convert)
  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 203, in _apply
    module._apply(fn)
  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 203, in _apply
    module._apply(fn)
  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 225, in _apply
    param_applied = fn(param)
  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 441, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
RuntimeError: CUDA error: out of memory

但是此时显存可以查看到还有很多空间，明显不是显存不够，那是哪里出问题了？

我试过更换其他同类型模型，完全没有问题，可以跑通，但是换回来就又是显存溢出，那么肯定不是模型超大的问题了（模型确实不大），一度怀疑是模型损坏，但是经过检查，可以正常进行模型推理，那也不是损坏问题了。
这时脑子灵光一闪，是不是因为多显卡的问题，于是乎，改了一行code：

model.load_state_dict(torch.load(args.model_path, map_location='cuda:0'))

你没有看错，是显卡的分配问题，训练时使用的是4张显卡，而在使用的时候有一张显卡内存是不够用的，所以导致模型加载失败，那么只要让一张显卡加载模型即可。