The problem here is that DistributedDataParallel performs gradient averaging across processes by hooking into the AccumulateGrad 15 function. This allows for performing averaging for the last most gradients while autograd is still running.
Would it be possible for you to first compute the initial loss, call autograd.backward instead of autograd.grad, and have it accumulate the first order gradients in the model parameters? Then you could detach those and compute something else before letting the optimizer do its thing. If not, then you’ll have to perform your own averaging, I think.