Kendall_Reid:

self.running_avg_mean = self.running_avg_mean + self.momentum * (batch_ch_mean.detach() - self.running_avg_mean)


Because in DP, the python module object is replicated to run on each GPU in a different thread. However, this setattr assigns the updated the buffer to the replica, which is lost right afterwards. Instead, inplace updates to the buffer works because buffers in the replica on the first GPU share memory with the original one.