解决多卡加载预训练模型时，卡0总会比其他卡多占用显存，多卡占用显存不均

出现问题：

再跑stdc官方给出代码时，使用多卡加载pretrained模型时，卡0上比卡123多三个进程，卡123都只有一个进程，中断训练后卡0全部进程都中断。即卡0在正常训练时比卡123多占用显存。（使用nvidia-smi命令查看）

这里发现问题存在于load模型的时候，直接load的话会导致参数加载到之前保存模型的device上（大部分情况下应该是只用cuda0去保存），这里可以将load函数加一个参数为map_location解决：

model_weights = torch.load(path, map_location='cpu')

改完之后发现仍旧存在这个问题

后来查看代码的时候发现，在网络的class内部，init_weight的时候也有一个load函数。这个函数内部无map_location参数。

因此，加入这个参数后如果卡0仍旧多占用显存和进程的时候，可能是某个隐藏的类别内部存在torch.load函数。多找找加上map_location参数即可使得全部的卡占用相同显存。