未验证 提交 0de1979d 编写于 作者: G guo ran 提交者: GitHub

fix input_op infer_nd_sbp (#6642)

Co-authored-by: Noneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>
上级 4bee2ddd
......@@ -70,10 +70,12 @@ Maybe<void> InputOp::InferNdSbpSignature(
std::function<Maybe<const NdSbpInferHint*>(const std::string&)> NdSbpInferHint4Ibn) const {
const auto& parallel_hierarchy = parallel_desc.hierarchy();
const InterfaceBlobConf& blob_conf = op_conf().input_conf().blob_conf();
cfg::NdSbp& tick_nd_sbp = (*nd_sbp_signature->mutable_bn_in_op2nd_sbp())["tick"];
tick_nd_sbp.clear_sbp_parallel();
FOR_RANGE(int64_t, i, 0, parallel_hierarchy->NumAxes()) {
tick_nd_sbp.mutable_sbp_parallel()->Add()->mutable_broadcast_parallel();
if (op_conf().input_conf().has_tick()) {
cfg::NdSbp& tick_nd_sbp = (*nd_sbp_signature->mutable_bn_in_op2nd_sbp())["tick"];
tick_nd_sbp.clear_sbp_parallel();
FOR_RANGE(int64_t, i, 0, parallel_hierarchy->NumAxes()) {
tick_nd_sbp.mutable_sbp_parallel()->Add()->mutable_broadcast_parallel();
}
}
cfg::NdSbp& out_nd_sbp = (*nd_sbp_signature->mutable_bn_in_op2nd_sbp())["out"];
JUST(InterfaceOpUtil::ParseNdSbpFromBlobConf(blob_conf, parallel_desc, &out_nd_sbp));
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册