diff --git a/python/paddle/distributed/fleet/data_generator/__init__.py b/python/paddle/distributed/fleet/data_generator/__init__.py
index 230ada2abec06292001965aa19e2e97c60d232b2..2288aca43f75149ce84742d9a5d954e28ce5f3b9 100644
--- a/python/paddle/distributed/fleet/data_generator/__init__.py
+++ b/python/paddle/distributed/fleet/data_generator/__init__.py
@@ -11,6 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 
-from .data_generator import DataGenerator  # noqa: F401
+from .data_generator import DataGenerator, MultiSlotDataGenerator  # noqa: F401
 
 __all__ = []