Browse Source

return actual dataset size when take_count is -1

pull/15122/head
ms_yan 5 years ago
parent
commit
5cc23859dd
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      mindspore/ccsrc/minddata/dataset/engine/ir/datasetops/take_node.cc

+ 1
- 1
mindspore/ccsrc/minddata/dataset/engine/ir/datasetops/take_node.cc View File

@@ -68,7 +68,7 @@ Status TakeNode::GetDatasetSize(const std::shared_ptr<DatasetSizeGetter> &size_g
}
int64_t num_rows;
RETURN_IF_NOT_OK(children_[0]->GetDatasetSize(size_getter, estimate, &num_rows));
*dataset_size = std::min(static_cast<int64_t>(take_count_), num_rows);
*dataset_size = take_count_ == -1 ? num_rows : std::min(static_cast<int64_t>(take_count_), num_rows);
dataset_size_ = *dataset_size;
return Status::OK();
}


Loading…
Cancel
Save