Can you try this? It’s not the most efficient but it should be better.
    def initialize(self, x):
        # compute batch statistics
        mean = x.mean(2).mean(0).resize(1, x.size(1), 1)
        mean_sq = (x**2).mean(2).mean(0).resize(1, x.size(1), 1)
        self.batch_size = x.size(0)
        assert x is not None
        assert mean is not None
        assert mean_sq is not None
        # build detached variables to avoid backprop to graph to compute mean and mean_sq
        # we will manually backprop those in hooks
        self.mean = autograd.Variable(mean.data.clone(), requires_grad = True)  # new code
        self.mean_sq = autograd.Variable(mean_sq.data.clone(), requires_grad = True)  # new code
        self.mean.register_hook(lambda grad: mean.backward(grad, retain_graph = True))  # new code
        self.mean_sq.register_hook(lambda grad: mean_sq.backward(grad, retain_graph = True))  # new code
        # compute reference output
        out = self._normalize(x, mean, mean_sq)
        self.reference_output = out.detach_()  # change, just to remove unnecessary saved graph
        return mean, mean_sq
   
    def get_ref_batch_stats(self):
        return self.mean, self.mean_sq

    def forward(self, x):
        if self.reference_output is None:
            ref_mean, ref_mean_sq = self.initialize(x)  # here
        else:  # here
            ref_mean, ref_mean_sq = self.get_ref_batch_stats()  # here
        new_coeff = 1. / (self.batch_size + 1.)
        old_coeff = 1. - new_coeff
        new_mean = x.mean(2).mean(0).resize_as()
        new_mean_sq = (x**2).mean(2).mean(0).resize_as(ref_mean_sq)  # here
        mean = new_coeff * new_mean + old_coeff * ref_mean  # here
        mean_sq = new_coeff * new_mean_sq + old_coeff * ref_mean_sq  # here
        x = self._normalize(x, mean, mean_sq)
        return x

    def _normalize(self, x, mean, mean_sq):
        assert self.eps is not None
        assert mean_sq is not None
        assert mean is not None
        assert len(x.size()) == 3
        gamma = Variable(torch.Tensor(1, x.size(1), 1).cuda().normal_(1., 0.02))
        beta = Variable(torch.cuda.FloatTensor(1, x.size(1), 1).fill_(0))
        std = torch.sqrt(self.eps + mean_sq - mean**2)
        x = x - mean
        x = x / std
        x = x * gamma
        x = x + beta
        return x

Additionally, you should add gamma and beta as parameters for them to be optimized. They are currently not.