code: def replace_linear(linear, mask):
      """
      Parameters
      ----------
      linear : torch.nn.Linear
          The linear module to be replace
      mask : ModuleMasks
          The masks of this module
  
      Returns
      -------
      torch.nn.Linear
          The new linear module
      """
      assert isinstance(mask, ModuleMasks)
      assert mask.input_mask is not None
      assert mask.output_mask is None
      assert not mask.param_masks
      index = mask.input_mask.mask_index[-1]
      in_features = index.size()[0]
      _logger.debug(
          "replace linear with new in_features: %d", in_features
      )
      new_linear = torch.nn.Linear(
          in_features=in_features,
          out_features=linear.out_features,
          bias=linear.bias is not None,
      )
      new_linear.to(linear.weight.device)
      new_linear.weight.data = torch.index_select(
          linear.weight.data, -1, index.to(linear.weight.device)
      )
      if linear.bias is not None:
          new_linear.bias.data.copy_(linear.bias.data)
      return new_linear
  
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 2 for calling a global function with example usage: replace_linear(linear, mask) and returns: new_linear