code: def conv2d_mask(module_masks, mask):
      """
      Infer input and output shape from weight mask
  
      Parameters
      ----------
      module_masks : ModuleMasks
          The ModuleMasks instance of the conv2d
      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
      """
  
      def convert_to_coarse_mask(mask, dim=0):
          """
          Parameters
          ----------
          mask : dict
              Weight mask from user provided mask file
          dim: int
              0: filter pruning
              1: channel pruning
  
          Returns
          -------
          LongTensor, CoarseMask, CoarseMask
              Index of the masked dimension, weight mask, bias mask
          """
          assert "weight" in mask
          assert isinstance(mask["weight"], torch.Tensor)
          assert dim in [0, 1]
  
          weight_mask = mask["weight"]
  
          sum_idx = (1, 2, 3) if dim == 0 else (0, 2, 3)
          index = torch.nonzero(
              weight_mask.abs().sum(sum_idx) != 0, as_tuple=True
          )[0]
  
          index = index.long().to(weight_mask.device)
          weight_cmask = CoarseMask(num_dim=4)
          weight_cmask.add_index_mask(dim=dim, index=index)
          bias_cmask = None
          if dim == 0 and "bias" in mask and mask["bias"] is not None:
              bias_index = torch.nonzero(mask["bias"], as_tuple=True)[0]
              assert torch.all(
                  torch.eq(index, bias_index)
              ), "bias mask should be consistent with weight mask"
              bias_cmask = CoarseMask(num_dim=1)
              bias_cmask.add_index_mask(dim=0, index=bias_index)
          return index, weight_cmask, bias_cmask
  
      index, weight_cmask, bias_cmask = convert_to_coarse_mask(
          mask, dim=conv_prune_dim
      )
  
      if index is None:
          # TODO: fine grained mask speedup
          return None, None
      # deal with coarse grain mask
      # mask conflict should be solved by fix_mask_conflict before speedup
      if "weight" in module_masks.param_masks:
          assert module_masks.param_masks["weight"] == weight_cmask
      else:
          module_masks.set_param_masks("weight", weight_cmask)
          if conv_prune_dim == 0:
              module_masks.set_param_masks("bias", bias_cmask)
  
      io_cmask = CoarseMask(num_dim=4)
      io_cmask.add_index_mask(dim=1, index=index)
  
      if conv_prune_dim == 0:
          if module_masks.output_mask is None:
              module_masks.set_output_mask(io_cmask)
          else:
              assert module_masks.output_mask == io_cmask
          return None, module_masks.output_mask
      else:
          if module_masks.input_mask is None:
              module_masks.set_input_mask(io_cmask)
          else:
              assert module_masks.input_mask == io_cmask
          return module_masks.input_mask, None
  
imports: import logging
  import torch
label_desc: Use image node_id 21 for calling a global function with example usage: conv2d_mask(module_masks, mask) and returns: index,  weight_cmask,  bias_cmask, None,  None, None,  module_masks, module_masks,  None