From 4f4857a4ac4d09483f72465e5adcd29f38b03b16 Mon Sep 17 00:00:00 2001
From: Billel Mokeddem <billel.mokeddem.ml@gmail.com>
Date: Fri, 15 Nov 2024 16:16:15 +0400
Subject: [PATCH] Fix: Change embeddings to embedding (#2738)

fix: change embeddings to embedding

Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-135.us-west-2.compute.internal>
---
 .../models/custom_modeling/mamba_modeling.py                    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/text_generation_server/models/custom_modeling/mamba_modeling.py b/server/text_generation_server/models/custom_modeling/mamba_modeling.py
index 07284e6a..5a9c0588 100644
--- a/server/text_generation_server/models/custom_modeling/mamba_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/mamba_modeling.py
@@ -212,7 +212,7 @@ class MambaModel(nn.Module):
         try:
             self.lm_head = SpeculativeHead.load(config, f"{prefix}.embeddings", weights)
         except RuntimeError:
-            self.lm_head = SpeculativeHead.load(config, f"{prefix}.embeddings", weights)
+            self.lm_head = SpeculativeHead.load(config, f"{prefix}.embedding", weights)
         self.config = config
 
     def forward(