ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(model): 修复加载状态字典时的键存在性检查

This commit is contained in:
ViperEkura 2025-11-25 21:03:10 +08:00
parent fb85aaf6a6
commit 5daf63a7a4
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
khaosz/model/transformer.py
View File

@ -83,8 +83,9 @@ class Transformer(nn.Module):
# same tensor
state_dict[lm_head_key] = state_dict[embed_key]
else:
# use clone to avoid sharing the same tensor
state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
if lm_head_key not in state_dict and embed_key in state_dict:
# use clone to avoid sharing the same tensor
state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
return super().load_state_dict(state_dict, strict, assign)