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:
177
177
:param module: module to check, may be offloaded
178
178
:return: onload device of module
179
179
"""
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
183
183
184
- param = next (module .parameters (recurse = False ), None )
184
+ param = next (submodule .parameters (recurse = False ), None )
185
185
if param is not None :
186
186
return param .device
187
187
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" )
191
189
return torch .device ("cpu" )
192
190
193
191
You can’t perform that action at this time.
0 commit comments