def save_registry():
    """Save the registry state and restore after the context manager exits."""
    saved_registry_state = copy.deepcopy(catalogue.REGISTRY)

    yield

    catalogue.REGISTRY = saved_registry_state
Save the registry state and restore after the context manager exits.
llmfoundry/utils/registry_utils.py