Commit 9199f074 authored by Ted Themistokleous's avatar Ted Themistokleous
Browse files

Fix unsqueeze on parse_if shape checks

parent 69977052
...@@ -98,7 +98,7 @@ struct parse_if : op_parser<parse_if> ...@@ -98,7 +98,7 @@ struct parse_if : op_parser<parse_if>
else_out_shapes.at(0).lens(), else_out_shapes.at(0).lens(),
else_out_shapes.at(0).strides()}; else_out_shapes.at(0).strides()};
auto reshape_ins = then_mdl->insert_instruction( auto reshape_ins = then_mdl->insert_instruction(
convert_ins, migraphx::make_op("unsqueeze", {{"axes", {0, 1}}}), convert_ins); convert_ins, migraphx::make_op("unsqueeze", {{"axes", {1}}}), convert_ins);
then_mdl->replace_return({reshape_ins}); then_mdl->replace_return({reshape_ins});
} }
...@@ -119,7 +119,7 @@ struct parse_if : op_parser<parse_if> ...@@ -119,7 +119,7 @@ struct parse_if : op_parser<parse_if>
then_out_shapes.at(0).lens(), then_out_shapes.at(0).lens(),
then_out_shapes.at(0).strides()}; then_out_shapes.at(0).strides()};
auto reshape_ins = then_mdl->insert_instruction( auto reshape_ins = then_mdl->insert_instruction(
convert_ins, migraphx::make_op("unsqueeze", {{"axes", {0, 1}}}), convert_ins); convert_ins, migraphx::make_op("unsqueeze", {{"axes", {1}}}), convert_ins);
else_mdl->replace_return({reshape_ins}); else_mdl->replace_return({reshape_ins});
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment