diff --git a/python/paddle/fluid/incubate/fleet/tests/ctr_dataset_reader.py b/python/paddle/fluid/incubate/fleet/tests/ctr_dataset_reader.py
index 32ba0e512f54ee142cc02cd6e6a36589daffdd4e..1407e92fd386b3312f4f6f695ecc5a36849ab401 100644
--- a/python/paddle/fluid/incubate/fleet/tests/ctr_dataset_reader.py
+++ b/python/paddle/fluid/incubate/fleet/tests/ctr_dataset_reader.py
@@ -19,7 +19,7 @@ import tarfile
 import os
 
 import paddle
-import paddle.fluid.incubate.data_generator as data_generator
+from paddle.distributed.fleet.dataset import data_generator as data_generator
 from paddle.fluid.log_helper import get_logger
 
 logger = get_logger(
diff --git a/python/paddle/fluid/tests/unittests/ctr_dataset_reader.py b/python/paddle/fluid/tests/unittests/ctr_dataset_reader.py
index 15e98481c26b20de4e9fa493fa022380ba1fcd63..6799d943df1bc6b0544ba724bc234721bf0b5ae6 100644
--- a/python/paddle/fluid/tests/unittests/ctr_dataset_reader.py
+++ b/python/paddle/fluid/tests/unittests/ctr_dataset_reader.py
@@ -22,8 +22,7 @@ import random
 import warnings
 
 import paddle
-import paddle.fluid.incubate.data_generator as data_generator
-
+from paddle.distributed.fleet.dataset import data_generator as data_generator
 logging.basicConfig()
 logger = logging.getLogger("paddle")
 logger.setLevel(logging.INFO)
diff --git a/python/paddle/fluid/tests/unittests/simnet_dataset_reader.py b/python/paddle/fluid/tests/unittests/simnet_dataset_reader.py
index 41eadc13a2ad26ac15b0623147dae5771f371a12..dea23ae97f3adb7b22601e42f67966d547a4ba20 100644
--- a/python/paddle/fluid/tests/unittests/simnet_dataset_reader.py
+++ b/python/paddle/fluid/tests/unittests/simnet_dataset_reader.py
@@ -21,7 +21,7 @@ import tarfile
 import random
 
 import paddle
-import paddle.fluid.incubate.data_generator as data_generator
+from paddle.distributed.fleet.dataset import data_generator as data_generator
 
 logging.basicConfig()
 logger = logging.getLogger("paddle")