diff --git a/tests/ut/python/dataset/test_paddeddataset.py b/tests/ut/python/dataset/test_paddeddataset.py
index afe52f3e858..4dbc187447b 100644
--- a/tests/ut/python/dataset/test_paddeddataset.py
+++ b/tests/ut/python/dataset/test_paddeddataset.py
@@ -504,7 +504,7 @@ def test_celeba_padded():
     count = 0
     for _ in data.create_dict_iterator(num_epochs=1, output_numpy=True):
         count = count + 1
-    assert count == 4 
+    assert count == 2
 
 
 if __name__ == '__main__':