move moe tensor

hpcaitech · Aug 14, 2023 · da1789f · da1789f
1 parent 00cb06e
commit da1789f
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/colossalai/nn/layer/moe/experts.py b/colossalai/nn/layer/moe/experts.py
@@ -8,7 +8,7 @@
 
 from colossalai.context import ParallelMode, seed
 from colossalai.context.moe_context import MOE_CONTEXT
-from colossalai.nn.layer.moe.moe_param import set_moe_param_info
+from colossalai.tensor.moe_tensor.api import set_moe_param_info
 from colossalai.utils import get_current_device
 from colossalai.zero.legacy.init_ctx import no_shard_zero_decrator
 

diff --git a/colossalai/nn/layer/moe/moe_param.py → colossalai/tensor/moe_tensor/api.py b/colossalai/nn/layer/moe/moe_param.py → colossalai/tensor/moe_tensor/api.py
diff --git a/colossalai/zero/low_level/low_level_optim.py b/colossalai/zero/low_level/low_level_optim.py
@@ -16,7 +16,7 @@
 )
 from colossalai.interface import OptimizerWrapper
 from colossalai.logging import get_dist_logger
-from colossalai.nn.layer.moe.moe_param import is_moe_param
+from colossalai.tensor.moe_tensor.api import is_moe_param
 # from colossalai.tensor import ColoParameter, ProcessGroup
 from colossalai.utils.cuda import get_current_device