def get_layer_submodules(module, submodule_types=None, layer_name_filter=None, name=""):
    """Get the submodules of a module based on the submodule types."""
    submodule_types = submodule_types or [torch.nn.Conv2d, torch.nn.Linear, transformers.Conv1D]
    if type(module) in submodule_types:
        if layer_name_filter and not any(s in name for s in layer_name_filter):
            # skip this layer
            return {}
        return {name: module}

    submodules = {}
    for submodule_name_k, submodule in module.named_children():
        submodule_name = name + "." + submodule_name_k if name else submodule_name_k
        submodules.update(get_layer_submodules(submodule, submodule_types, layer_name_filter, submodule_name))
    return submodules
Get the submodules of a module based on the submodule types.
olive/passes/pytorch/sparsegpt_utils.py