@smth
Admittedly, still not getting it. I  tried creating a mean only bn module
def __init__(self, num_features,momentum=0.1):
    super(_MeanOnlyBatchNorm, self).__init__()
    self.num_features = num_features
    self.momentum = momentum
    self.register_buffer('running_mean', torch.zeros(num_features))
    self.reset_parameters()

def forward(self, input):
    mu = torch.mean(input,dim=0).data
    self.running_mean = self.momentum*mu + (1-self.momentum) * self.running_mean
    return torch.add(input, -self.running_mean)

but get errors. I’m lost.
EDIT:
I’ve settled on this, but still am not sure if it’s correct.
def forward(self, input):

    # mu = Variable(torch.mean(input,dim=0, keepdim=True).data, requires_grad=False)
    if self.training is True:
        mu = input.mean(dim=0, keepdim=True)
        mu = self.momentum * mu + (1 - self.momentum) * Variable(self.running_mean)

        self.running_mean = mu.data
        return input.add_(-mu)
    else:
        return  input.add_(-Variable(self.running_mean))