use apply rather than recursion

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
neuralmagic · Dec 19, 2024 · df3e186 · df3e186
1 parent 38d7dbf
commit df3e186
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/src/compressed_tensors/utils/offload.py b/src/compressed_tensors/utils/offload.py
@@ -235,7 +235,7 @@ def delete_offload_parameter(module: torch.nn.Module, name: str):
 
 @check_accelerate(fallback=contextlib.nullcontext())
 @contextlib.contextmanager
-def disable_hf_hook(module: torch.nn.Module, recurse: bool = False):
+def disable_hf_hook(module: torch.nn.Module):
     hooks = {}
 
     def collect_hooks(module):
@@ -244,10 +244,7 @@ def collect_hooks(module):
             hooks[module] = module._hf_hook
             remove_hook_from_module(module)
 
-        for submodule in module.children():
-            collect_hooks(submodule)
-
-    collect_hooks(module)
+    module.apply(collect_hooks)
 
     yield