diff --git a/internlm/model/__init__.py b/internlm/model/__init__.py
index b0fe77d..d1dc634 100644
--- a/internlm/model/__init__.py
+++ b/internlm/model/__init__.py
@@ -5,12 +5,14 @@ from .embedding import Embedding1D, RotaryEmbedding
 from .linear import FeedForward, RewardModelLinear, ScaleColumnParallelLinear
 from .metrics import AccPerplex
 from .modeling_internlm import build_model_with_cfg
+from .moe import MoE
 from .multi_head_attention import MHA
 from .utils import gather_forward_split_backward
 
 __all__ = [
     "Embedding1D",
     "FeedForward",
+    "MoE",
     "RotaryEmbedding",
     "RewardModelLinear",
     "ScaleColumnParallelLinear",