experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Also check dst_format

Browse Source
This commit is contained in:
Kaixi Hou 2020-09-23 12:04:10 -07:00
parent d602375436
commit 9b947dd637
1 changed files with 2 additions and 1 deletions

3
tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.cc
View File

@ -1379,7 +1379,8 @@ Status ReduceTransposer::TransposeNode(TransposeContext* context,
std::string src_format = context->src_format;
std::string dst_format = context->dst_format;
// Update the format from 4D to 5D layout if necessary.
bool allow_5d = rank == 5 && (src_format == "NHWC" || src_format == "NCHW");
bool allow_5d = rank == 5 && (src_format == "NHWC" || src_format == "NCHW") &&
(dst_format == "NHWC" || dst_format == "NCHW");
if (allow_5d) {
std::string src_format_3d = src_format == "NHWC" ? "NDHWC" : "NCDHW";
std::string dst_format_3d = dst_format == "NHWC" ? "NDHWC" : "NCDHW";