Fix new shape in tflite reshape importer

kendryte · May 11, 2020 · bb231ed · bb231ed
1 parent ec8b8cb
commit bb231ed
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/importer/tflite/ops/reshape.cpp b/src/importer/tflite/ops/reshape.cpp
@@ -23,8 +23,9 @@ DEFINE_TFLITE_LOWER(RESHAPE)
 {
     auto &input = get_tensor(op.inputs(), 0);
     auto &options = *op.builtin_options_as_ReshapeOptions();
+    auto new_shape = load_axis<int32_t>(get_tensor(op.inputs(), 1));
 
-    auto node = graph_.emplace<reshape>(to_data_type(input.type()), get_shape(input.shape()), get_axis(*options.new_shape()));
+    auto node = graph_.emplace<reshape>(to_data_type(input.type()), get_shape(input.shape()), new_shape);
 
     input_tensors_.emplace(&node->input(), op.inputs()->Get(0));
     output_tensors_.emplace(op.outputs()->Get(0), &node->output());