From 5f5d02fbf15a4a4e998a6855675e2abb81dc85a5 Mon Sep 17 00:00:00 2001
From: Sayak Paul <spsayakpaul@gmail.com>
Date: Mon, 26 May 2025 20:15:29 -0700
Subject: [PATCH] Make group offloading compatible with torch.compile()
 (#11605)

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 7a8970aeed..565f8f1ff8 100644
--- a/src/diffusers/hooks/group_offloading.py
+++ b/src/diffusers/hooks/group_offloading.py
@@ -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."""