Skip to content

Commit b78eb8f

Browse files
authored
fully remove dispatch (#427)
Signed-off-by: Kyle Sayers <[email protected]>
1 parent a2bfc03 commit b78eb8f

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

src/compressed_tensors/utils/offload.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -579,6 +579,7 @@ def remove_dispatch(module: torch.nn.Module) -> torch.nn.Module:
579579
remove_hook_from_module(module, recurse=True)
580580
if hasattr(module, "hf_device_map"):
581581
delattr(module, "hf_device_map")
582+
module.to("cpu")
582583

583584
return module
584585

0 commit comments

Comments
 (0)