class_docstr: A BatchNorm layer used in several flows
class_docstr_tok: ["A","BatchNorm","layer","used","in","several","flows"]
code: def forward(self, x):
      if self.training:
          self.batch_mean = x.mean(0).data
          self.batch_var = x.var(0).data
  
          self.running_mean.mul_(1 - self.momentum).add_(
              self.batch_mean * self.momentum
          )
          self.running_var.mul_(1 - self.momentum).add_(
              self.batch_var * self.momentum
          )
  
          mean = self.batch_mean
          var = self.batch_var
  
      else:
          mean = self.running_mean
          var = self.running_var
  
      y = (
          (x - mean) / (var + self.eps).sqrt()
      ) * self.log_gamma.exp() + self.beta
  
      log_abs_det_jac = self.log_gamma - 0.5 * (var + self.eps).log()
  
      output = ModelOutput(
          out=y,
          log_abs_det_jac=log_abs_det_jac.expand_as(x).sum(dim=-1),
      )
  
      return output
  
file_path: benchmark_VAE/src/pythae/models/normalizing_flows/layers.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from pythae.models.base.base_utils import ModelOutput
label_desc: Use image node_id 2 for calling the BatchNorm obj's underlying member method code with example usage: obj.forward(x) and returns: output