Pytorch:多块GPU调用细节问题及Pytorch的nn.DataParallel解释_userwarning: was asked to gather along dimension 0-CSDN博客

网站介绍:文章浏览阅读7k次,点赞11次,收藏23次。我们用实验室带有多块卡的GPU服务器,当我们在上面跑程序的时候,当迭代次数或者epoch足够大的时候,我们通常会使用nn.DataParallel函数加入以下代码段来用多个GPU来加速训练。device_ids = [0, 1]net = torch.nn.DataParallel(net, device_ids=device_ids)然鹅,也常常会由于正在运行程序的0卡显存不够放不下我们新来的模型初始化数据而导致显存OOM的问题不断,使得模型无法训练,同时报出警告:UserWarning_userwarning: was asked to gather along dimension 0, but all input tensors we