See https://el-hult.github.io/2021/02/25/invertible-normalizer.html

batchnorm1d_nf.py

class MemorizingNormalizer(nn.Module):
    def __init__(self, d, eps, rho):
        super().__init__()
        self.means = nn.Parameter(torch.zeros(d), requires_grad=False)
        self.vars = nn.Parameter(torch.ones(d), requires_grad=False)
        self.eps = nn.Parameter(torch.tensor(eps, dtype=float), requires_grad=False)
        self.rho = nn.Parameter(torch.tensor(rho, dtype=float), requires_grad=False)
        
    def forward(self, x):
        self.means.data = self.means * self.rho + (1 - self.rho) * x.mean(axis=0)
        self.vars.data = self.vars * self.rho + (1 - self.rho) * x.var(axis=0)
        varse = self.vars + self.eps
        lj = -0.5 * torch.log(torch.prod(varse)) # log jacobian determinant
        z = (x - self.means) / torch.sqrt(varse)
        return z, lj
      
    def inverse(self, z):
        return z * torch.sqrt(self.vars + self.eps) + self.means