From 37add20783c2b2e693f75cf3da16afb5c48b01e8 Mon Sep 17 00:00:00 2001 From: hesham Date: Mon, 8 Feb 2021 20:39:23 -0500 Subject: [PATCH] Fix caching problem after cleanup --- .../api/python/bindings/dataset/include/datasets_bindings.cc | 5 ++--- mindspore/dataset/engine/datasets.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/mindspore/ccsrc/minddata/dataset/api/python/bindings/dataset/include/datasets_bindings.cc b/mindspore/ccsrc/minddata/dataset/api/python/bindings/dataset/include/datasets_bindings.cc index eb8ba07c09e..976ba6848da 100644 --- a/mindspore/ccsrc/minddata/dataset/api/python/bindings/dataset/include/datasets_bindings.cc +++ b/mindspore/ccsrc/minddata/dataset/api/python/bindings/dataset/include/datasets_bindings.cc @@ -87,9 +87,8 @@ PYBIND_REGISTER(DatasetNode, 1, ([](const py::module *m) { return num_workers ? self->SetNumWorkers(*num_workers) : self; }) .def("set_cache_client", - [](std::shared_ptr self) { - std::shared_ptr dc = nullptr; - return self->SetDatasetCache(dc); + [](std::shared_ptr self, std::shared_ptr cc) { + return self->SetDatasetCache(toDatasetCache(std::move(cc))); }) .def( "Zip", diff --git a/mindspore/dataset/engine/datasets.py b/mindspore/dataset/engine/datasets.py index c944c8824f6..971799b544d 100644 --- a/mindspore/dataset/engine/datasets.py +++ b/mindspore/dataset/engine/datasets.py @@ -3475,7 +3475,7 @@ class GeneratorDataset(MappableDataset): new_op.sampler = None new_op.sample_fn = sample_fn new_op.source_len = min(new_op.source_len, - new_op.num_samples) if new_op.num_samples is not None else new_op.source_len + new_op.num_samples) if new_op.num_samples != 0 else new_op.source_len iter(self.source) except TypeError: # Use generator function if input callable