fix bug

PaddlePaddle · Aug 16, 2024 · e59a8b7 · e59a8b7
1 parent 877e7c4
commit e59a8b7
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/paddlenlp/trainer/auto_trainer.py b/paddlenlp/trainer/auto_trainer.py
@@ -727,9 +727,8 @@ def _load_from_checkpoint(self, resume_from_checkpoint=None):
             if self.args.to_static:
                 parameter_to_structured_name = self.model_wrapped._parameter_to_structured_name
             else:
-                for state_dict_name, sub_state_dict in state_dict.items():
-                    for state_name, state_value in sub_state_dict.items():
-                        parameter_to_structured_name[state_value.name] = state_name
+                for state_name, state_value in self.model_wrapped.state_dict().items():
+                    parameter_to_structured_name[state_value.name] = state_name
 
             if self.args.resume_form_hybrid_parallel:
                 CheckpointConverter(