diff --git a/mindspeed_llm/mindspore/mindspore_adaptor.py b/mindspeed_llm/mindspore/mindspore_adaptor.py
index 467a40c5c5905da5c30c23a188eb24725133c77c..27066e0e4a3795bd31a53ed3bad33ccd6fc60827 100644
--- a/mindspeed_llm/mindspore/mindspore_adaptor.py
+++ b/mindspeed_llm/mindspore/mindspore_adaptor.py
@@ -323,8 +323,8 @@ class MindSporeAdaptation(MegatronAdaptationABC):
                         zerc_alltoall_token_unperm2)
 
         from mindspeed.mindspore.core.pipeline_parallel.fb_overlap.modules.token_dispatcher import PackProb, UnpackProb
-        MindSporeAdaptation.register('mindspeed.mindspore.core.pipeline_parallel.fb_overlap.modules.token_dispatcher.PackProb', PackProb)
-        MindSporeAdaptation.register('mindspeed.mindspore.core.pipeline_parallel.fb_overlap.modules.token_dispatcher.UnpackProb', UnpackProb)
+        MindSporeAdaptation.register('mindspeed.core.pipeline_parallel.fb_overlap.modules.token_dispatcher.PackProb', PackProb)
+        MindSporeAdaptation.register('mindspeed.core.pipeline_parallel.fb_overlap.modules.token_dispatcher.UnpackProb', UnpackProb)
 
         if args.gradient_accumulation_fusion:
             from mindspeed.mindspore.ops.npu_matmul_add import npu_matmul_add_fp32