def get_atomic_layers(
    module: _nn.Module,
    layer_types: _Union[_List[str], _List[_Type]],
    name_prefix: str = "",
) -> _Dict[str, _nn.Module]:
    """
    Returns a dictionary of layer_name: layer for every layer in the module which
    matches the types specified in layers_to_find.
    """
    if isinstance(module, tuple(layer_types)):
        return {name_prefix: module}
    result = {}
    for name, child in module.named_children():
        result.update(
            get_atomic_layers(
                child,
                layer_types=layer_types,
                name_prefix=name_prefix + "." + name if name_prefix != "" else name,
            )
        )

    return result

    Returns a dictionary of layer_name: layer for every layer in the module which
    matches the types specified in layers_to_find.
    
coremltools/optimize/torch/_utils/torch_utils.py