fix #5295

Former-commit-id: c76873b0eb8225f6e6bfc7223c6012387dceb8ed
2024-08-29 20:30:18 +08:00
parent 77341ee3c4
commit 7c6785d3df
4 changed files with 7 additions and 8 deletions
--- a/src/llamafactory/train/callbacks.py
+++ b/src/llamafactory/train/callbacks.py
@@ -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