pytorch多卡分布式训练卡住的问题

最近需要模型进行多任务学习，完成代码后单卡跑是没有问题的。但是多卡跑基本上会出现卡在第一个迭代的情况。忙了一天总算是解决了。总结一下我这里的原因：

我的代码里执行了torch.distributed.all_reduce() 的操作，但这个操作只服务于一个task，导致的结果是多卡训练中，部分卡执行了torch.distributed.all_reduce()，部分卡没有执行该操作。 知道原因，解答方案也比较明了，要么让所有卡都执行torch.distributed.all_reduce()，要么让所有卡都不执行该操作。 我的模型不适应前者，因而我是把这部分给删了，从而解决。

实际上，尝试了很多种方案，最后发现问题出现在这里。这部分内容具体参考如下
https://github.com/pytorch/pytorch/issues/20630#issuecomment-730184591