From 1986116527b10ae348711be2b6a7daff136adf82 Mon Sep 17 00:00:00 2001
From: zhanglei <ryancheung98@163.com>
Date: Mon, 25 Sep 2023 15:36:35 +0800
Subject: [PATCH] add moe module to `__init__.py`

---
 internlm/model/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

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",