1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00

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

wip: check if we can make go compile compat
This commit is contained in:
Sayak Paul
2025-05-26 20:15:29 -07:00
committed by GitHub
parent 53748217e6
commit 5f5d02fbf1

View File

@@ -113,6 +113,7 @@ class ModuleGroup:
finally:
pinned_dict = None
@torch.compiler.disable()
def onload_(self):
r"""Onloads the group of modules to the onload_device."""
torch_accelerator_module = (
@@ -165,6 +166,7 @@ class ModuleGroup:
if self.record_stream:
buffer.data.record_stream(current_stream)
@torch.compiler.disable()
def offload_(self):
r"""Offloads the group of modules to the offload_device."""