From 5cc23859dd7bf9fa01d5711d2fef0d83d514269a Mon Sep 17 00:00:00 2001 From: ms_yan Date: Tue, 13 Apr 2021 22:01:06 +0800 Subject: [PATCH] return actual dataset size when take_count is -1 --- .../ccsrc/minddata/dataset/engine/ir/datasetops/take_node.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mindspore/ccsrc/minddata/dataset/engine/ir/datasetops/take_node.cc b/mindspore/ccsrc/minddata/dataset/engine/ir/datasetops/take_node.cc index ebc0809d08..fde5894f25 100644 --- a/mindspore/ccsrc/minddata/dataset/engine/ir/datasetops/take_node.cc +++ b/mindspore/ccsrc/minddata/dataset/engine/ir/datasetops/take_node.cc @@ -68,7 +68,7 @@ Status TakeNode::GetDatasetSize(const std::shared_ptr &size_g } int64_t num_rows; RETURN_IF_NOT_OK(children_[0]->GetDatasetSize(size_getter, estimate, &num_rows)); - *dataset_size = std::min(static_cast(take_count_), num_rows); + *dataset_size = take_count_ == -1 ? num_rows : std::min(static_cast(take_count_), num_rows); dataset_size_ = *dataset_size; return Status::OK(); }