I think the problem might be coming from this line: insightface/partial_fc.py at master · deepinsight/insightface · GitHub 22. dist.all_gather doesn’t perform autograd recording, so from an autograd point of view features is not used to produce loss and this might be causing the issue. One way to validate this would be to initialize total_features with copies of features (so its recorded as part of autograd) and see if that resolves the issue?