Reuse module in sys.modules if exists

huggingface · Apr 25, 2024 · 6985bee · 6985bee
1 parent 1fccb99
commit 6985bee
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/transformers/dynamic_module_utils.py b/src/transformers/dynamic_module_utils.py
@@ -200,11 +200,13 @@ def get_class_in_module(class_name: str, module_path: Union[str, os.PathLike]) -
     name = os.path.normpath(module_path).rstrip(".py").replace(os.path.sep, ".")
     module_path = str(Path(HF_MODULES_CACHE) / module_path)
     module_spec = importlib.util.spec_from_file_location(name, location=module_path)
-    module = sys.modules.pop(name, None)
+    module = sys.modules.get(name)
     if module is None:
         module = importlib.util.module_from_spec(module_spec)
+        # insert it into sys.modules before any loading begins
+        sys.modules[name] = module
+    # reload in both cases
     module_spec.loader.exec_module(module)
-    sys.modules[name] = module
     return getattr(module, class_name)