code: def compute_BatchNorm2d_flops(module, inp, out):
      assert isinstance(module, nn.BatchNorm2d)
      assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())
      in_c, in_h, in_w = inp.size()[1:]
      batch_flops = np.prod(inp.shape)
      if module.affine:
          batch_flops *= 2
      return batch_flops
  
file_path: tensorwatch/tensorwatch/model_graph/torchstat/compute_flops.py
imports: import torch.nn
  import torch
  import numpy
  import math
label_desc: Use image node_id 4 for calling a global function with example usage: compute_BatchNorm2d_flops(module, inp, out) and returns: batch_flops