expose

Signed-off-by: Mayank Mishra <[email protected]>
IBM · Oct 7, 2024 · 8459daf · 8459daf
1 parent 4c3c815
commit 8459daf
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/dolomite_engine/hf_models/models/moe_dolomite_TP/layer.py b/dolomite_engine/hf_models/models/moe_dolomite_TP/layer.py
@@ -7,7 +7,7 @@
 from ..gpt_dolomite_TP.layer import MLP_TP
 from ..moe_dolomite import MoEDolomiteConfig
 from ..moe_dolomite.layer import SparseMoEBlock
-from .moe_TP.scatter import ScatterMoE_TP
+from .moe_TP import ScatterMoE_TP
 
 
 class SparseMoEBlock_TP(SparseMoEBlock):

diff --git a/dolomite_engine/hf_models/models/moe_dolomite_TP/moe_TP/__init__.py b/dolomite_engine/hf_models/models/moe_dolomite_TP/moe_TP/__init__.py
@@ -0,0 +1 @@
+from .scatter import ScatterMoE_TP