Solved by deJQK in post #2 
              
              
                Following this discussion (Implementing Batchnorm in Pytorch. Problem with updating self.running_mean and self.running_var - #2 by SeoHyeong), it is implemented as the following: 
class SepBN2d(nn.BatchNorm2d):
    def __init__(self, num_features, eps=1e-5, momentum=0.1,
                 affine=True…