Unverified Commit 81df9c85 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

Unwrap models everywhere (#5789)

more debug
parent bfe94a39
......@@ -560,10 +560,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
register_dict = {name: (None, None)}
else:
# register the config from the original module, not the dynamo compiled one
if is_compiled_module(module):
not_compiled_module = module._orig_mod
else:
not_compiled_module = module
not_compiled_module = _unwrap_model(module)
library = not_compiled_module.__module__.split(".")[0]
......@@ -666,7 +663,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
# Dynamo wraps the original model in a private class.
# I didn't find a public API to get the original class.
if is_compiled_module(sub_model):
sub_model = sub_model._orig_mod
sub_model = _unwrap_model(sub_model)
model_cls = sub_model.__class__
save_method_name = None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment