Skip to content

Commit 0754e8f

Browse files
Unwrap models everywhere (huggingface#5789)
more debug
1 parent 869218c commit 0754e8f

File tree

1 file changed

+2
-5
lines changed

1 file changed

+2
-5
lines changed

pipelines/pipeline_utils.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -560,10 +560,7 @@ def register_modules(self, **kwargs):
560560
register_dict = {name: (None, None)}
561561
else:
562562
# register the config from the original module, not the dynamo compiled one
563-
if is_compiled_module(module):
564-
not_compiled_module = module._orig_mod
565-
else:
566-
not_compiled_module = module
563+
not_compiled_module = _unwrap_model(module)
567564

568565
library = not_compiled_module.__module__.split(".")[0]
569566

@@ -666,7 +663,7 @@ def is_saveable_module(name, value):
666663
# Dynamo wraps the original model in a private class.
667664
# I didn't find a public API to get the original class.
668665
if is_compiled_module(sub_model):
669-
sub_model = sub_model._orig_mod
666+
sub_model = _unwrap_model(sub_model)
670667
model_cls = sub_model.__class__
671668

672669
save_method_name = None

0 commit comments

Comments
 (0)