Implementing Batchnorm in Pytorch. Problem with updating self.running_mean and self.running_var

I’m not sure, why the running stats updates are not gathered to the default device, but using

self.running_mean.copy_(...)
# instead of
self.running_mean = (...)

seem to perform the updates properly.

CC @SeoHyeong

2 Likes