File tree Expand file tree Collapse file tree 1 file changed +2
-5
lines changed Expand file tree Collapse file tree 1 file changed +2
-5
lines changed Original file line number Diff line number Diff line change @@ -560,10 +560,7 @@ def register_modules(self, **kwargs):
560
560
register_dict = {name : (None , None )}
561
561
else :
562
562
# 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 )
567
564
568
565
library = not_compiled_module .__module__ .split ("." )[0 ]
569
566
@@ -666,7 +663,7 @@ def is_saveable_module(name, value):
666
663
# Dynamo wraps the original model in a private class.
667
664
# I didn't find a public API to get the original class.
668
665
if is_compiled_module (sub_model ):
669
- sub_model = sub_model . _orig_mod
666
+ sub_model = _unwrap_model ( sub_model )
670
667
model_cls = sub_model .__class__
671
668
672
669
save_method_name = None
You can’t perform that action at this time.
0 commit comments