code: def replace_linear(linear, masks):
      """
      This function will replace the original linear according to
      the infered masks. This function support the fine-grained and
      coarse-grained sparsity. In the fine-grained scenario, this function
      will remove the whole column/row that happen to be totally covered by
      the masks.
  
      Parameters
      ----------
      linear : torch.nn.Linear
          The linear module to be replace
      masks : Tuple of the input masks, output masks and weight masks
          Tuple of the masks, for example
          ([input_m1, input_m2], [output_m], {'weight':weight_m})
  
      Returns
      -------
      torch.nn.Linear
          The new linear module
      """
      in_masks, output_mask, weight_mask = masks
      assert isinstance(linear, nn.Linear)
      if len(in_masks) != 1:
          raise InputsNumberError()
      if not isinstance(output_mask, torch.Tensor):
          raise OutputTypeError(type(output_mask), torch.Tensor)
  
      in_mask = in_masks[0]
  
      weight_mask = weight_mask["weight"]
      # N C K
      pruned_in, remained_in = convert_to_coarse_mask(in_mask, -1)
      pruned_out, remained_out = convert_to_coarse_mask(output_mask, -1)
      n_remained_in = weight_mask.size(1) - pruned_in.size(0)
      n_remained_out = weight_mask.size(0) - pruned_out.size(0)
      remained_in, remained_out = remained_in.to(
          linear.weight.device
      ), remained_out.to(linear.weight.device)
      _logger.info(
          "replace linear with new in_features: %d, out_features: %d",
          n_remained_in,
          n_remained_out,
      )
      need_bias = False
      if linear.bias is not None:
          need_bias = True
      new_linear = torch.nn.Linear(
          in_features=n_remained_in,
          out_features=n_remained_out,
          bias=need_bias,
      )
      new_linear.to(linear.weight.device)
      # Copy the remained weight from the original module
      with torch.no_grad():
          tmp_weight_data = torch.index_select(
              linear.weight.data, 0, remained_out
          )
          new_linear.weight.data = torch.index_select(
              tmp_weight_data, 1, remained_in
          )
  
          if linear.bias is not None:
              new_linear.bias.data = torch.index_select(
                  linear.bias.data, 0, remained_out
              )
  
      return new_linear
  
file_path: nni/nni/compression/speedup/replacement.py
imports: import logging
  import math
  import torch
  import torch.nn
  from .error_code import EmptyLayerError, ShapeMisMatchError, InputsNumberError, OutputTypeError, UnBalancedGroupError
label_desc: Use image node_id 5 for calling a global function with example usage: replace_linear(linear, masks) and returns: new_linear