fix loading from pretrained for sharded model with `torch_dtype="auto" (#18061)

83d2d745 · Nouamane Tazi · GitHub · 7996ef74 · 83d2d745
Unverified Commit 83d2d745 authored Jul 27, 2022 by Nouamane Tazi Committed by GitHub Jul 27, 2022
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

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

No files found.
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -2073,7 +2073,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
                        elif not is_sharded:
                            torch_dtype = get_state_dict_dtype(state_dict)
                        else:
-                            one_state_dict = load_state_dict(resolved_archive_file)
+                            one_state_dict = load_state_dict(resolved_archive_file[0])
                            torch_dtype = get_state_dict_dtype(one_state_dict)
                            del one_state_dict  # free CPU memory
                    else: