code: def detect_channel_prune_type(
      graph_module: torch.fx.GraphModule,
      masks: Dict[str, Dict[str, torch.Tensor]],
  ) -> str:
      """
      User can prune a channel through two ways: 1) prune
      the corresponding filter of the conv layer(all the
      filter related pruner), 2) prune the BN layers that
      followed after a conv(Slim pruner). This function find
      the pruning type of the masks.
  
      Parameters
      ----------
      graph_module: torch.fx.GraphModule
          GraphModule object which the mask can be applied on.
      masks: dict
          A dict object that stores the masks.
  
      Returns:
      -------
      prune_type: str
          Could be Filter or BatchNorm
      """
      for layer_name in masks:
          sub_module = graph_module.get_submodule(layer_name)
          if sub_module is None or (
              not isinstance(sub_module, torch.nn.BatchNorm2d)
          ):
              return "Filter"
      # if all masks are for batchnorm layers, then the prune_type is BatchNorm
      # Note, actually we currently do not support pruning both Conv and BatchNorm
      # at the same time.
      return "BatchNorm"
  
file_path: nni/nni/compression/speedup/mask_conflict.py
imports: import logging
  from typing import Dict
  import numpy
  import torch
  import torch.fx
  from .dependency import build_channel_dependency, build_group_dependency, build_weight_sharing_dependency
label_desc: Use image node_id 4 for calling a global function with example usage: detect_channel_prune_type(graph_module, masks) and returns: str, str