Thank you! With the code below I get the following error:
RuntimeError: Trying to backward through the graph a second time, 
but the buffers have already been freed. Specify retain_graph=True 
when calling backward the first time.     

    def initialize(self, x):
        # compute batch statistics
        self.mean = x.mean(2).mean(0).resize(1, x.size(1), 1)
        self.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 self.mean is not None
        assert self.mean_sq is not None
        # compute reference output
        out = self._normalize(x, self.mean, self.mean_sq)
        self.reference_output = out

    def forward(self, x):
        if self.reference_output is None:
            self.initialize(x)
        new_coeff = 1. / (self.batch_size + 1.)
        old_coeff = 1. - new_coeff
        new_mean = x.mean(2).mean(0).resize_as(self.mean)
        new_mean_sq = (x**2).mean(2).mean(0).resize_as(self.mean_sq)
        mean = new_coeff * new_mean + old_coeff * self.mean
        mean_sq = new_coeff * new_mean_sq + old_coeff * self.mean_sq
        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