Update callbacks.py

This commit is contained in:
hoshi-hiyouga 2024-07-10 13:32:20 +08:00 committed by GitHub
parent cff89a2e89
commit 39cd89ce17
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions

View File

@ -79,7 +79,7 @@ def fix_valuehead_checkpoint(
if name.startswith("v_head."): if name.startswith("v_head."):
v_head_state_dict[name] = param v_head_state_dict[name] = param
else: else:
decoder_state_dict[name.replace("pretrained_model.", "",1)] = param decoder_state_dict[name.replace("pretrained_model.", "", 1)] = param
model.pretrained_model.save_pretrained( model.pretrained_model.save_pretrained(
output_dir, state_dict=decoder_state_dict or None, safe_serialization=safe_serialization output_dir, state_dict=decoder_state_dict or None, safe_serialization=safe_serialization