diff --git a/python/setup.py.in b/python/setup.py.in
index 773166400347ab550f82e4fabcb0d89b90818fc2..7c87cacbe62b79d4f383bd867c3d4f2901387ef2 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -150,6 +150,7 @@ packages=['paddle',
           'paddle.distributed.fleet.meta_optimizers',
           'paddle.distributed.fleet.runtime',
           'paddle.distributed.fleet.dataset',
+          'paddle.distributed.fleet.data_generator',
           'paddle.distributed.fleet.metrics',
           'paddle.distributed.fleet.proto',
           'paddle.distributed.fleet.utils',
diff --git a/setup.py b/setup.py
index af558c2ef0b42b68e47fe98ebd626c9b9034bef9..edccee638b6c9aa9b45cb19fb451d93e518f8483 100644
--- a/setup.py
+++ b/setup.py
@@ -184,6 +184,7 @@ packages = [
     'paddle.distributed.fleet.meta_optimizers',
     'paddle.distributed.fleet.runtime',
     'paddle.distributed.fleet.dataset',
+    'paddle.distributed.fleet.data_generator',
     'paddle.distributed.fleet.metrics',
     'paddle.distributed.fleet.proto',
     'paddle.distributed.fleet.utils',