diff --git a/dygraph/core/train.py b/dygraph/core/train.py
index 6c64532fffc0951e79596cea17971b8d51e534f6..59df4fd595e658409d84854a1d28045f6e4297d3 100644
--- a/dygraph/core/train.py
+++ b/dygraph/core/train.py
@@ -17,7 +17,8 @@ import os
 import paddle.fluid as fluid
 from paddle.fluid.dygraph.parallel import ParallelEnv
 from paddle.fluid.io import DataLoader
-from paddle.incubate.hapi.distributed import DistributedBatchSampler
+# from paddle.incubate.hapi.distributed import DistributedBatchSampler
+from paddle.io import DistributedBatchSampler
 
 import dygraph.utils.logger as logger
 from dygraph.utils import load_pretrained_model