Update modeling_Llamoe.py

modeling_Llamoe.py

@@ -745,7 +745,7 @@ class LlamoeDecoderLayer(nn.Module):
 
         self.self_attn = LLAMOE_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
 
-        self.block_sparse_moe = LlamoeBlockSparseTop2MLPSparseMoeBlock(config)
+        self.block_sparse_moe = LlamoeBlockSparseTop2MLP(config)
         self.input_layernorm = LlamoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = LlamoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)