From fe8089871e1227fd8ab3c1ea0997a9fb91098050 Mon Sep 17 00:00:00 2001
From: Pierrick HYMBERT <pierrick.hymbert@gmail.com>
Date: Sat, 6 Apr 2024 23:27:29 +0200
Subject: [PATCH] model: dbrx: fix missing embedding tensor, mix with output
 layer

---
 gguf-py/gguf/tensor_mapping.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py
index 0b455d113..0bf7ae7b2 100644
--- a/gguf-py/gguf/tensor_mapping.py
+++ b/gguf-py/gguf/tensor_mapping.py
@@ -24,6 +24,7 @@ class TensorNameMap:
             "backbone.embedding",                        # mamba
             "backbone.embeddings",                       # mamba-hf
             "transformer.in_out_embed",                  # Grok
+            "transformer.wte.weight",                    # dbrx
         ),
 
         # Token type embeddings
@@ -52,7 +53,7 @@ class TensorNameMap:
             "output",                    # llama-pth bloom internlm2
             "word_embeddings_for_head",  # persimmon
             "lm_head.linear",            # phi2
-            "transformer.wte.weight",    # dbrx
+            "lm_head.weight",            # dbrx
         ),
 
         # Output norm