def dmas_bn(module, input, output):
    """DMAs estimation for `torch.nn.modules.batchnorm._BatchNorm`"""
    input_dma = input.numel()
    ops_dma = module.running_mean.numel() + module.running_var.numel() + 1
    if module.affine:
        ops_dma += module.weight.data.numel() + module.bias.data.numel()
    if module.momentum:
        ops_dma += 1
    if module.training and module.track_running_stats:
        ops_dma += module.num_batches_tracked.numel(
            ) + module.running_mean.numel() + module.running_var.numel()
    output_dma = output.numel()
    return input_dma + ops_dma + output_dma
