update new strategy for UnsortedSegments to fit stra check

4 years ago · e2d713ddbf
parent 2b8737a908
commit e2d713ddbf
1 changed files with 3 additions and 0 deletions
--- a/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc
+++ b/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc
@ -370,6 +370,9 @@ Strategys MakeRecSearchStrategy(const std::shared_ptr<Graph> &graph,
  if (iter_ops >= ops.size()) {
    MS_LOG(EXCEPTION) << "Failure: Operators' elements out of range.";
  }
+  if (graph->nodes[iter_graph].apply.op_type == kRecUnsortedSegmentOp) {
+    return MakeDataParallelStrategy(graph, ops, iter_graph, iter_ops);
+  }

  StrategyPtr origin_strategy = ops[iter_ops]->strategy();
  Strategys strategies;