code: def batchnorm2d_mask(module_masks, mask):
      """
      Infer input and output shape from weight mask
      Parameters
      ----------
      module_masks : ModuleMasks
          The ModuleMasks instance of the batchnorm2d
      mask : dict
          The mask of its weights, from the user provided mask file
      Returns
      -------
      CoarseMask, CoarseMask
          The mask of its input tensor, the mask of its output tensor
      """
      assert "weight" in mask and "bias" in mask
      sum_mask = mask["weight"] + mask["bias"]
      nonzero_index = torch.nonzero(sum_mask, as_tuple=True)[0]
      # infer shape of parameters
      param_cmask = CoarseMask(num_dim=1)
      param_cmask.add_index_mask(dim=0, index=nonzero_index)
      module_masks.set_param_masks("weight", param_cmask)
      module_masks.set_param_masks("bias", param_cmask)
      # infer shape of input tensor
      input_cmask = CoarseMask(num_dim=4)
      input_cmask.add_index_mask(
          dim=1, index=torch.nonzero(mask["weight"], as_tuple=True)[0]
      )
      module_masks.set_input_mask(input_cmask)
      # infer shape of output tensor
      output_cmask = CoarseMask(num_dim=4)
      output_cmask.add_index_mask(dim=1, index=nonzero_index)
      module_masks.set_output_mask(output_cmask)
      return input_cmask, output_cmask
  
imports: import logging
  import torch
label_desc: Use image node_id 19 for calling a global function with example usage: batchnorm2d_mask(module_masks, mask) and returns: input_cmask,  output_cmask