Skip to content

GitLab

Unverified Commit 5f5d02fb authored by Sayak Paul's avatar Sayak Paul Committed by GitHub
Browse files

Make group offloading compatible with torch.compile() (#11605) 

wip: check if we can make go compile compat
parent 53748217
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
src/diffusers/hooks/group_offloading.py
View file @ 5f5d02fb
...@@ -113,6 +113,7 @@ class ModuleGroup: ...@@ -113,6 +113,7 @@ class ModuleGroup:
finally: finally:
pinned_dict = None pinned_dict = None
@torch.compiler.disable()
def onload_(self): def onload_(self):
r"""Onloads the group of modules to the onload_device.""" r"""Onloads the group of modules to the onload_device."""
torch_accelerator_module = ( torch_accelerator_module = (
...@@ -165,6 +166,7 @@ class ModuleGroup: ...@@ -165,6 +166,7 @@ class ModuleGroup:
if self.record_stream: if self.record_stream:
buffer.data.record_stream(current_stream) buffer.data.record_stream(current_stream)
@torch.compiler.disable()
def offload_(self): def offload_(self):
r"""Offloads the group of modules to the offload_device.""" r"""Offloads the group of modules to the offload_device."""
......
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