github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Fixes bug by raising exception on size mismatch

pull/5493/head
Jianbing Wu 2024-03-22 19:52:40 +08:00 committed by GitHub
parent 6df844b8c4
commit 3596bd60e7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/checkpoint_io/utils.py
View File

@ -558,6 +558,10 @@ def load_state_dict_into_model(
missing_keys = missing_keys.append(sub_missing_keys) missing_keys = missing_keys.append(sub_missing_keys)
if strict: if strict:
if len(error_msgs) > 0:
raise RuntimeError(
"Error(s) in loading state_dict for {}:\n\t{}".format(model.__class__.__name__, "\n\t".join(error_msgs))
)
if len(unexpected_keys) > 0: if len(unexpected_keys) > 0:
error_msgs = "Unexpected key(s) in state_dict: {}. ".format( error_msgs = "Unexpected key(s) in state_dict: {}. ".format(
", ".join('"{}"'.format(k) for k in unexpected_keys) ", ".join('"{}"'.format(k) for k in unexpected_keys)