utils/get_layers.py

import torch

def get_layers(model,layer_types = (torch.nn.Linear, torch.nn.Conv2d,torch.nn.BatchNorm2d)):
  """
  Extract all layers from a PyTorch model dynamically.

  Args:
      model (torch.nn.Module): The PyTorch model.
      layer_types (tuple): Add the layers you want to track

  Returns:
      List[Tuple[str, nn.Module]]: A list of layer names and their corresponding modules.
  """

  return [(name, layer) for name , layer in model.named_modules() if isinstance(layer, layer_types)]