Batch Normalization Inverse Computation using Pytorch

I have implemented a BatchNorm class for computing both BatchNormalization and its inverse, when I test it with tensors with 1 batch, it works properly, but when I test it for multi-batch tensors it doesn't work right. Code:

class BatchNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.zeros(1, dim), requires_grad=True)
        self.beta = nn.Parameter(torch.zeros(1, dim), requires_grad=True)
        self.batch_mean = None
        self.batch_var = None

    def forward(self, x, reverse=False):
        B, C, W, H = x.shape
        if(reverse == True):
            return self.reverse(x)
        if self.training:
            if(B>1):
                m = x.mean(dim=0)
                v = x.var(dim=0) + self.eps  # torch.mean((x - m) ** 2, axis=0) + self.eps
            else:
                m = torch.zeros(C, W, H)
                v = torch.zeros(C, W, H) + self.eps
            self.batch_mean = None
        else:
            if self.batch_mean is None:
                self.set_batch_stats_func(x)
            m = self.batch_mean.clone()
            v = self.batch_var.clone()
        B, C, W, H = x.shape
        gamma = self.gamma.unsqueeze(2).unsqueeze(3)
        gamma = torch.repeat_interleave(gamma, H, dim=2)
        gamma = torch.repeat_interleave(gamma, W, dim=3)
        beta = self.beta.unsqueeze(2).unsqueeze(3)
        beta = torch.repeat_interleave(beta, H, dim=2)
        beta = torch.repeat_interleave(beta, W, dim=3)
        #print('x_hat:', x_hat)
        x_hat = (x - m) / torch.sqrt(v)
        x_hat = x_hat * torch.exp(gamma) + beta
        
        x_2 = (x_hat - beta) * torch.exp(-gamma) * torch.sqrt(v) + m
        #print('forward: dist:', torch.dist(x, x_2))
        #print('forward: x:', x[0,0,:3,:3])
        #print('forward: x_2:', x_2[0,0,:3,:3])
        #print('forward: x_hat:', x_hat[0,0,:3,:3])
        
        log_det = torch.sum(gamma - 0.5 * torch.log(v))
        return x_hat, log_det

    def reverse(self, x):
        B, C, W, H = x.shape
        if self.training:
            if(B>1):
                m = x.mean(dim=0)
                v = x.var(dim=0) + self.eps  # torch.mean((x - m) ** 2, axis=0) + self.eps
            else:
                m = torch.zeros(C, W, H)
                v = torch.zeros(C, W, H) + self.eps
            self.batch_mean = None
        else:
            if self.batch_mean is None:
                self.set_batch_stats_func(x)
            m = self.batch_mean
            v = self.batch_var

        B, C, W, H = x.shape
        gamma = self.gamma.unsqueeze(2).unsqueeze(3)
        gamma = torch.repeat_interleave(gamma, H, dim=2)
        gamma = torch.repeat_interleave(gamma, W, dim=3)
        beta = self.beta.unsqueeze(2).unsqueeze(3)
        beta = torch.repeat_interleave(beta, H, dim=2)
        beta = torch.repeat_interleave(beta, W, dim=3)
        x_hat = (x - beta) * torch.exp(-gamma) * torch.sqrt(v) + m
        #print('reverse: dist:', torch.dist(x, x_hat))
        #print('reverse: x:', x[0,0,:3,:3])
        #print('reverse: x_hat:', x_hat[0,0,:3,:3])
        log_det = torch.sum(-gamma + 0.5 * torch.log(v))
        return x_hat, log_det

    def set_batch_stats_func(self, x):
        print("setting batch stats for validation")
        self.batch_mean = x.mean(dim=0)
        self.batch_var = x.var(dim=0) + self.eps

Test on one batch tensor:

x = torch.rand(1,10,100,100)
Batch = BatchNorm(10)
x1,_ = Batch(x, False)
x2,_ = Batch(x1, True)
torch.dist(x,x2)

and the output is about zero, it means that both forward and backward paths are working properly, but for multi-batch tensors:

x = torch.rand(3,10,100,100)
Batch = BatchNorm(10)
x1,_ = Batch(x, False)
x2,_ = Batch(x1, True)
torch.dist(x,x2)

In this case, the result(difference between input and reconstructed input) is a huge number. However, it must be near zero.