From 74146053facc705a4f85eb303675ee26e18271f7 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 23 May 2025 15:48:17 +0530 Subject: [PATCH] wip: check if we can make go compile compat --- src/diffusers/hooks/group_offloading.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 7a8970aeedb2..565f8f1ff860 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -113,6 +113,7 @@ def _pinned_memory_tensors(self): 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 @@ def onload_(self): 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."""