423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-16 20:00:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix #5295

Browse Source
This commit is contained in:
hiyouga
2024-08-29 20:30:18 +08:00
parent 364b757e30
commit ad72f3e065
4 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/llamafactory/train/callbacks.py
View File

@@ -75,8 +75,7 @@ def fix_valuehead_checkpoint(
state_dict: Dict[str, torch.Tensor] = torch.load(path_to_checkpoint, map_location="cpu")
os.remove(path_to_checkpoint)
decoder_state_dict = {}
v_head_state_dict = {}
decoder_state_dict, v_head_state_dict = {}, {}
for name, param in state_dict.items():
if name.startswith("v_head."):
v_head_state_dict[name] = param