Unverified Commit 4b03b502 authored by liangel-02's avatar liangel-02 Committed by GitHub
Browse files

update torchao safetensors impl (#30155)


Signed-off-by: default avatarAngel Li <liangel@meta.com>
parent 4c6fd258
...@@ -641,7 +641,6 @@ def safetensors_weights_iterator( ...@@ -641,7 +641,6 @@ def safetensors_weights_iterator(
if safetensors_load_strategy == "eager": if safetensors_load_strategy == "eager":
loading_desc += " (eager)" loading_desc += " (eager)"
state_dict = {}
leftover_state_dict: dict[str, torch.Tensor] = {} leftover_state_dict: dict[str, torch.Tensor] = {}
for st_file in tqdm( for st_file in tqdm(
...@@ -667,6 +666,7 @@ def safetensors_weights_iterator( ...@@ -667,6 +666,7 @@ def safetensors_weights_iterator(
) )
with safe_open(st_file, framework="pt") as f: with safe_open(st_file, framework="pt") as f:
state_dict = {}
for name in f.keys(): # noqa: SIM118 for name in f.keys(): # noqa: SIM118
state_dict[name] = f.get_tensor(name) state_dict[name] = f.get_tensor(name)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment