Skip to content

Commit 38cbdd1

Browse files
authored
[Accelerate] Fix typos in get_execution_device (#365)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
1 parent 10e4e55 commit 38cbdd1

File tree

1 file changed

+5
-7
lines changed

1 file changed

+5
-7
lines changed

src/compressed_tensors/utils/offload.py

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff 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

0 commit comments

Comments
 (0)