Skip to content

GitLab

Unverified Commit e2033d2d authored by Dhruv Nair's avatar Dhruv Nair Committed by GitHub
Browse files

Fix model offload bug when key isn't present (#5030) 

* fix model offload bug when key isn't present

* make style
parent 19edca82
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
src/diffusers/pipelines/pipeline_utils.py
View file @ e2033d2d
...@@ -1252,7 +1252,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): ...@@ -1252,7 +1252,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
self._all_hooks = [] self._all_hooks = []
hook = None hook = None
for model_str in self.model_cpu_offload_seq.split("->"): for model_str in self.model_cpu_offload_seq.split("->"):
model = all_model_components.pop(model_str) model = all_model_components.pop(model_str, None)
if not isinstance(model, torch.nn.Module): if not isinstance(model, torch.nn.Module):
continue continue
......
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