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