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."""