[BugFix] Fix a bug in loading safetensors (#1732)

f5a37c6c · Woosuk Kwon · GitHub · 32c927b5 · f5a37c6c
Unverified Commit f5a37c6c authored Nov 20, 2023 by Woosuk Kwon Committed by GitHub Nov 20, 2023
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

vllm/model_executor/weight_utils.py vllm/model_executor/weight_utils.py +1 -1

No files found.
--- a/vllm/model_executor/weight_utils.py
+++ b/vllm/model_executor/weight_utils.py
@@ -240,7 +240,7 @@ def hf_model_weights_iterator(
    elif use_safetensors:
        for st_file in hf_weights_files:
            with safe_open(st_file, framework="pt") as f:
-                for name in f:
+                for name in f.keys():  # noqa: SIM118
                    param = f.get_tensor(name)
                    yield name, param
    else: