From f7329b1a0e49893beba5bcbc8cedc2eb53fc89a3 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Thu, 18 Jan 2024 01:12:31 +0800 Subject: [PATCH] Update llamafy_internlm2.py Former-commit-id: 3ca5915a4fcd3d28d10a47bf9f2188b5cf8393a8 --- tests/llamafy_internlm2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/llamafy_internlm2.py b/tests/llamafy_internlm2.py index 41b12512..3cd59f96 100644 --- a/tests/llamafy_internlm2.py +++ b/tests/llamafy_internlm2.py @@ -40,7 +40,7 @@ def save_weight( llama2_state_dict: Dict[str, torch.Tensor] = OrderedDict() for key, value in tqdm(internlm2_state_dict.items(), desc="Convert format"): if "output" in key: - llama2_state_dict["lm_head"] = value + llama2_state_dict[key.replace("output", "lm_head")] = value elif "tok_embeddings" in key: llama2_state_dict[key.replace("tok_embeddings", "embed_tokens")] = value elif "attention_norm" in key: