Extremely experimental fix!

huggingface · Aug 23, 2024 · e9e154b · e9e154b
1 parent 0a7af19
commit e9e154b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/models/auto/auto_factory.py b/src/transformers/models/auto/auto_factory.py
@@ -429,6 +429,7 @@ def from_config(cls, config, **kwargs):
             model_class = get_class_from_dynamic_module(class_ref, repo_id, **kwargs)
             if os.path.isdir(config._name_or_path):
                 model_class.register_for_auto_class(cls.__name__)
+                cls.register(config.__class__, model_class, exist_ok=True)
             else:
                 cls.register(config.__class__, model_class, exist_ok=True)
             _ = kwargs.pop("code_revision", None)
@@ -554,6 +555,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
             _ = hub_kwargs.pop("code_revision", None)
             if os.path.isdir(pretrained_model_name_or_path):
                 model_class.register_for_auto_class(cls.__name__)
+                cls.register(config.__class__, model_class, exist_ok=True)
             else:
                 cls.register(config.__class__, model_class, exist_ok=True)
             return model_class.from_pretrained(