From 37ffca4df8648b03a3dee63fe60aead0e425b620 Mon Sep 17 00:00:00 2001 From: limingqi107 Date: Thu, 24 Dec 2020 12:37:11 +0800 Subject: [PATCH] ascend support ps cache sink_size as 1 --- mindspore/train/dataset_helper.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mindspore/train/dataset_helper.py b/mindspore/train/dataset_helper.py index b561c722e5..8ba7eb9de2 100644 --- a/mindspore/train/dataset_helper.py +++ b/mindspore/train/dataset_helper.py @@ -99,7 +99,8 @@ def connect_network_with_dataset(network, dataset_helper): if (hasattr(dataset_iter, "sink_size") and dataset_iter.sink_size == 1) \ and (hasattr(dataset_iter, "sink_count") and dataset_iter.sink_count == 1) \ and context.get_context("device_target") == "Ascend" \ - and context.get_context("mode") == context.GRAPH_MODE: + and context.get_context("mode") == context.GRAPH_MODE \ + and ms_role != "MS_WORKER": if not hasattr(dataset, '__network__'): dataset.__network__ = network