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)]