lucas
/
LLaMA-Factory
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix #5305

This commit is contained in:
hiyouga 2024-08-29 20:16:01 +08:00
parent 0f5a0f64f7
commit 364b757e30
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -74,6 +74,7 @@ def fix_valuehead_checkpoint(
path_to_checkpoint = os.path.join(output_dir, WEIGHTS_NAME)
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 = {}
for name, param in state_dict.items():
@ -91,7 +92,6 @@ def fix_valuehead_checkpoint(
else:
torch.save(v_head_state_dict, os.path.join(output_dir, V_HEAD_WEIGHTS_NAME))
os.remove(path_to_checkpoint)
logger.info("Value head model saved at: {}".format(output_dir))