code: def get_output_batch_dims(t: Tensor, module: Module):
      if isinstance(module, (torch.nn.Linear, torch.nn.Bilinear)):
          batch_nums = math.prod(t.shape[:-1])
          batch_dims = [_ for _ in range(len(t.shape[:-1]))]
      elif isinstance(
          module, (torch.nn.Conv1d, torch.nn.ConvTranspose1d)
      ):
          batch_nums = math.prod(t.shape[:-2])
          batch_dims = [_ for _ in range(len(t.shape[:-2]))]
      elif isinstance(
          module, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)
      ):
          batch_nums = math.prod(t.shape[:-3])
          batch_dims = [_ for _ in range(len(t.shape[:-3]))]
      elif isinstance(
          module, (torch.nn.Conv3d, torch.nn.ConvTranspose3d)
      ):
          batch_nums = math.prod(t.shape[:-4])
          batch_dims = [_ for _ in range(len(t.shape[:-4]))]
      else:
          raise TypeError(
              f"Found unsupported module type in activation based pruner: {module.__class__.__name__}"
          )
      return batch_dims, batch_nums
  
file_path: nni/nni/compression/utils/mask_counter.py
imports: from copy import deepcopy
  import math
  from typing import Dict, List, Tuple
  import torch
  from torch import Tensor
  from torch.nn import Module
label_desc: Use image node_id 5 for calling a global function with example usage: get_output_batch_dims(t, module) and returns: batch_dims,  batch_nums