Update llamafy_internlm2.py

This commit is contained in:
hiyouga
2024-01-18 01:00:16 +08:00
parent c84a387c2c
commit 484becae1b

View File

@@ -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: