fix generatorDataset get_dataset_size

This commit is contained in:
yanghaitao 2020-05-29 14:48:37 +08:00
parent ad279e90fd
commit 7f5db74457
1 changed files with 2 additions and 1 deletions

View File

@ -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.