pc

llmfoundry/utils/registry_utils.py

@@ -1,10 +1,11 @@
 # Copyright 2024 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
+import copy
 import functools
 import importlib.util
-from contextlib import contextmanager
 import os
+from contextlib import contextmanager
 from pathlib import Path
 from types import ModuleType
 from typing import (Any, Callable, Dict, Generic, Optional, Sequence, Type,
@@ -176,11 +177,12 @@ def import_file(loc: Union[str, Path]) -> ModuleType:
         raise RuntimeError(f'Error executing {loc}') from e
     return module
 
+
 @contextmanager
 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
+    catalogue.REGISTRY = saved_registry_state

tests/fixtures/autouse.py

@@ -15,11 +15,13 @@
 REPO_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
 sys.path.append(REPO_DIR)
 
+
 @pytest.fixture(autouse=True)
 def save_registry_fixture():
     with save_registry():
         yield
 
+
 @pytest.fixture(autouse=True)
 def initialize_dist(request: pytest.FixtureRequest):
     """Initialize the default PyTorch distributed process group for tests."""