fix CSVDataset get_dataset_size interface

4 years ago · 853d10cd7d
parent beb86391fe
commit 853d10cd7d
1 changed files with 1 additions and 1 deletions
--- a/mindspore/dataset/engine/datasets.py
+++ b/mindspore/dataset/engine/datasets.py
@ -5488,7 +5488,7 @@ class CSVDataset(SourceDataset):
            num_rows = CsvOp.get_num_rows(self.dataset_files, self.column_names is None)
            self.dataset_size = get_num_rows(num_rows, self.num_shards)
            if self.num_samples is not None and self.num_samples < self.dataset_size:
-                self.dataset_size = num_rows
+                self.dataset_size = self.num_samples
        return self.dataset_size

    def is_shuffled(self):