code: def replace_batchnorm2d(norm, mask):
      """
      Parameters
      ----------
      norm : torch.nn.BatchNorm2d
          The batchnorm module to be replace
      mask : ModuleMasks
          The masks of this module
  
      Returns
      -------
      torch.nn.BatchNorm2d
          The new batchnorm module
      """
      assert isinstance(mask, ModuleMasks)
      assert "weight" in mask.param_masks and "bias" in mask.param_masks
      index = mask.param_masks["weight"].mask_index[0]
      num_features = index.size()[0]
      _logger.debug(
          "replace batchnorm2d with num_features: %d", num_features
      )
      new_norm = torch.nn.BatchNorm2d(
          num_features=num_features,
          eps=norm.eps,
          momentum=norm.momentum,
          affine=norm.affine,
          track_running_stats=norm.track_running_stats,
      )
      # assign weights
      new_norm.weight.data = torch.index_select(
          norm.weight.data, 0, index
      )
      new_norm.bias.data = torch.index_select(norm.bias.data, 0, index)
      if norm.track_running_stats:
          new_norm.running_mean.data = torch.index_select(
              norm.running_mean.data, 0, index
          )
          new_norm.running_var.data = torch.index_select(
              norm.running_var.data, 0, index
          )
      return new_norm
  
file_path: auptimizer/src/aup/compression/torch/speedup/compress_modules.py
imports: import logging
  import torch
  from .infer_shape import ModuleMasks
label_desc: Use image node_id 3 for calling a global function with example usage: replace_batchnorm2d(norm, mask) and returns: new_norm