code: def measure_layer_for_pruning(wrapper, x):
      def get_layer_type(layer):
          layer_str = str(layer)
          return layer_str[: layer_str.find("(")].strip()
  
      def get_layer_param(model):
          import operator
          import functools
  
          return sum(
              [
                  functools.reduce(operator.mul, i.size(), 1)
                  for i in model.parameters()
              ]
          )
  
      multi_add = 1
      layer = wrapper.module
      type_name = get_layer_type(layer)
  
      # ops_conv
      if type_name in ["Conv2d"]:
          out_h = int(
              (
                  x.size()[2]
                  + 2 * layer.padding[0]
                  - layer.kernel_size[0]
              )
              / layer.stride[0]
              + 1
          )
          out_w = int(
              (
                  x.size()[3]
                  + 2 * layer.padding[1]
                  - layer.kernel_size[1]
              )
              / layer.stride[1]
              + 1
          )
          wrapper.flops = (
              layer.in_channels
              * layer.out_channels
              * layer.kernel_size[0]
              * layer.kernel_size[1]
              * out_h
              * out_w
              / layer.groups
              * multi_add
          )
          wrapper.params = get_layer_param(layer)
      # ops_linear
      elif type_name in ["Linear"]:
          weight_ops = layer.weight.numel() * multi_add
          bias_ops = layer.bias.numel()
          wrapper.flops = weight_ops + bias_ops
          wrapper.params = get_layer_param(layer)
      return
  
file_path: auptimizer/src/aup/compression/torch/pruning/amc/lib/utils.py
imports: import os
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: measure_layer_for_pruning(wrapper, x) and returns: layer_str, sum