def get_named_leaf_modules(module: Module) -> Dict[str, Module]:
    """Returns all leaf modules of the model with their names."""
    return {
        name: mod for name, mod in module.named_modules() if not any(mod.children())
    }
Returns all leaf modules of the model with their names.
lightly/models/utils.py