From 7f5db74457e56d1e253b70a63501faa2b5bf4f4e Mon Sep 17 00:00:00 2001
From: yanghaitao <yanghaitao1@huawei.com>
Date: Fri, 29 May 2020 14:48:37 +0800
Subject: [PATCH] fix generatorDataset get_dataset_size

---
 mindspore/dataset/engine/datasets.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mindspore/dataset/engine/datasets.py b/mindspore/dataset/engine/datasets.py
index f3703b3850..f4e9679830 100644
--- a/mindspore/dataset/engine/datasets.py
+++ b/mindspore/dataset/engine/datasets.py
@@ -2948,7 +2948,8 @@ class GeneratorDataset(MappableDataset):
 
         if rows_from_sampler is None:
             return self._dataset_size
-
+        if self._dataset_size is None:
+            return None
         return min(rows_from_sampler, self._dataset_size)
 
     # manually set dataset_size as a temporary solution.