Fix unsqueeze on parse_if shape checks

9199f074 · Ted Themistokleous · 69977052 · 9199f074
Commit 9199f074 authored Aug 22, 2022 by Ted Themistokleous
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/onnx/parse_if.cpp src/onnx/parse_if.cpp +2 -2

No files found.
--- a/src/onnx/parse_if.cpp
+++ b/src/onnx/parse_if.cpp
@@ -98,7 +98,7 @@ struct parse_if : op_parser<parse_if>
                                  else_out_shapes.at(0).lens(),
                                  else_out_shapes.at(0).strides()};
                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});
            }
@@ -119,7 +119,7 @@ struct parse_if : op_parser<parse_if>
                                  then_out_shapes.at(0).lens(),
                                  then_out_shapes.at(0).strides()};
                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});
            }