don't load state_dict twice when using low_cpu_mem_usage in from_pretrained (#16602)

47c5c059 · Suraj Patil · GitHub · a2b7d19b · 47c5c059
Unverified Commit 47c5c059 authored Apr 06, 2022 by Suraj Patil Committed by GitHub Apr 06, 2022
Show whitespace changes
Inline Side-by-side

Showing with 0 additions and 1 deletion

src/transformers/modeling_utils.py src/transformers/modeling_utils.py +0 -1

No files found.
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -1823,7 +1823,6 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
                if is_sharded:
                    loaded_state_dict_keys = sharded_metadata["all_checkpoint_keys"]
                else:
-                    state_dict = load_state_dict(resolved_archive_file)
                    loaded_state_dict_keys = [k for k in state_dict.keys()]
                    del state_dict  # free CPU memory - will reload again later