Update llamafy_internlm2.py

Former-commit-id: 69b3cb768eda57b63f47cd35e5da3a59b57b7853
2024-01-18 01:00:16 +08:00
parent 636d8a886c
commit 48eb07c956
1 changed files with 1 additions and 1 deletions
--- a/tests/llamafy_internlm2.py
+++ b/tests/llamafy_internlm2.py
@@ -42,7 +42,7 @@ def save_weight(
        if "output" in key:
            llama2_state_dict["lm_head"] = value
        elif "tok_embeddings" in key:
-            llama2_state_dict["embed_tokens"] = value
+            llama2_state_dict[key.replace("tok_embeddings", "embed_tokens")] = value
        elif "attention_norm" in key:
            llama2_state_dict[key.replace("attention_norm", "input_layernorm")] = value
        elif "wqkv" in key: