diff --git a/python/paddle/fluid/tests/unittests/test_data_generator.py b/python/paddle/fluid/tests/unittests/test_data_generator.py
index acff800a17722ceee59837f7e0b01716b91901c8..310913f695b34bf59bcaa57e0b1f8e6ef12d9429 100644
--- a/python/paddle/fluid/tests/unittests/test_data_generator.py
+++ b/python/paddle/fluid/tests/unittests/test_data_generator.py
@@ -14,6 +14,7 @@ import paddle
 import unittest
 import paddle.distributed.fleet as fleet
 import os
+import sys
 
 
 class MyMultiSlotDataGenerator(fleet.MultiSlotDataGenerator):