fix typo of olmo

mlc-ai · Dec 17, 2024 · 1f3004a · 1f3004a
1 parent 286c33e
commit 1f3004a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/mlc_llm/model/olmo/olmo_model.py b/python/mlc_llm/model/olmo/olmo_model.py
@@ -102,7 +102,7 @@ def __post_init__(self):  # pylint: disable=too-many-branches
                 raise ValueError(f"'clip_qkv'({self.clip_qkv}) should be non-negative")
 
 
-class OLMoEebedding(nn.Embedding):
+class OLMoEmbedding(nn.Embedding):
     """The embedding module that can be shared with the final lm_head. From Qwen2Embedding."""
 
     def lm_head_forward(self, x: nn.Tensor):
@@ -248,7 +248,7 @@ def forward(  # pylint: disable=missing-function-docstring
 class OLMoModel(nn.Module):  # pylint: disable=missing-class-docstring
     def __init__(self, config: OLMoConfig):
         assert config.hidden_size % config.num_attention_heads == 0
-        self.embed_tokens = OLMoEebedding(config.vocab_size, config.hidden_size)
+        self.embed_tokens = OLMoEmbedding(config.vocab_size, config.hidden_size)
         self.layers = nn.ModuleList(
             [OLMoDecoderLayer(config) for _ in range(config.num_hidden_layers)]
         )