File tree Expand file tree Collapse file tree 1 file changed +5
-7
lines changed
src/compressed_tensors/utils Expand file tree Collapse file tree 1 file changed +5
-7
lines changed Original file line number Diff line number Diff line change @@ -177,17 +177,15 @@ def get_execution_device(module: torch.nn.Module) -> torch.device:
177177 :param module: module to check, may be offloaded
178178 :return: onload device of module
179179 """
180- for module in module .modules ():
181- if has_offloaded_params (module ):
182- return module ._hf_hook .execution_device
180+ for submodule in module .modules ():
181+ if has_offloaded_params (submodule ):
182+ return submodule ._hf_hook .execution_device
183183
184- param = next (module .parameters (recurse = False ), None )
184+ param = next (submodule .parameters (recurse = False ), None )
185185 if param is not None :
186186 return param .device
187187
188- warnings .warn (
189- f"Unable able to get execution device of { module } , falling back to CPU"
190- )
188+ warnings .warn (f"Unable to get execution device of { module } , falling back to CPU" )
191189 return torch .device ("cpu" )
192190
193191
You can’t perform that action at this time.
0 commit comments