Update modeling_Llamoe.py

modeling_Llamoe.py

@@ -755,7 +755,7 @@ LLAMA_ATTENTION_CLASSES = {
 
 
 class LlamoeDecoderLayer(nn.Module):
-    def __init__(self, config: GemmoeConfig, layer_idx: int):
+    def __init__(self, config: LlamoeConfig, layer_idx: int):
         super().__init__()
         self.hidden_size = config.hidden_size
 
@@ -860,7 +860,7 @@ LLAMA_START_DOCSTRING = r"""
 )
 
 class LlammoePreTrainedModel(PreTrainedModel):
-    config_class = GemmoeConfig
+    config_class = LlamoeConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _keep_in_fp32_modules = ["inv_freq", "rotary_emb", "cos_cached", "sin_cached"]
@@ -977,7 +977,7 @@ class LlamoeModel(GemmoePreTrainedModel):
         config: GemmoeConfig
     """
 
-    def __init__(self, config: GemmoeConfig):
+    def __init__(self, config: LlamoeConfig):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size