!22 fix the bug in broadcast_rewrite.

Merge pull request !22 from wYann/fix_broadcast_rewrite

!22 fix the bug in broadcast_rewrite.
Merge pull request !22 from wYann/fix_broadcast_rewrite
2be38102 · mindspore-ci-bot · Gitee · 4452f2eb · edfd65aa · 2be38102
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

src/pass/broadcast_rewrite.cc src/pass/broadcast_rewrite.cc +1 -1

未找到文件。
--- a/src/pass/broadcast_rewrite.cc
+++ b/src/pass/broadcast_rewrite.cc
@@ -80,7 +80,7 @@ class BroadcastVecRewriter : public IRMutator {
            }
          }

-          if (secDimFit && dimFit && forLoopFit) {
+          if (secDimFit && dimFit && forLoopFit && remainDim > 1) {
            newExtent = GetInt32Const(GetItem(forInfo.ops_, idx).as<For>()->extent) / remainDim;
            tmpBuffer = VarExpr("tmp_broadcast_" + std::to_string(broadBufferCount++) + "_local_UB", dtype);
            varList = dstInfo->var_;