ascend support ps cache sink_size as 1

4 years ago · 37ffca4df8
parent b42b9b6c80
commit 37ffca4df8
1 changed files with 2 additions and 1 deletions
--- 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