Skip to content

Commit 09ef275

Browse files
authored
【PIR】Change layout of binary infermeta (#72213)
* change layout of binary infermeta * fix ci
1 parent a037cd7 commit 09ef275

File tree

1 file changed

+9
-1
lines changed

1 file changed

+9
-1
lines changed

paddle/phi/infermeta/binary.cc

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1740,7 +1740,15 @@ void ElementwiseRawInferMeta(const MetaTensor& x,
17401740
promote_result = x.dtype();
17411741
}
17421742
out->set_dtype(promote_result);
1743-
out->set_layout(x.layout());
1743+
1744+
// layout need change when meet input layout contain kNHWC
1745+
auto layout = [&]() {
1746+
if (x.layout() == DataLayout::kNHWC || y.layout() == DataLayout::kNHWC)
1747+
return DataLayout::kNHWC;
1748+
return x.layout();
1749+
}();
1750+
1751+
out->set_layout(layout);
17441752
out->share_lod(x);
17451753
}
17461754

0 commit comments

Comments
 (0)