!12295 Fix Dataset cache problem after cleanup

From: @hfarahat
Reviewed-by: @robingrosman,@nsyca
Signed-off-by: @robingrosman
This commit is contained in:
mindspore-ci-bot 2021-02-10 02:57:49 +08:00 committed by Gitee
commit 8d55cd9d89
2 changed files with 3 additions and 4 deletions

View File

@ -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<DatasetNode> self) {
std::shared_ptr<DatasetCache> dc = nullptr;
return self->SetDatasetCache(dc);
[](std::shared_ptr<DatasetNode> self, std::shared_ptr<CacheClient> cc) {
return self->SetDatasetCache(toDatasetCache(std::move(cc)));
})
.def(
"Zip",

View File

@ -3472,7 +3472,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