From 38ef8d46847dfa90fc6f8135aca3ad30414b7c5d Mon Sep 17 00:00:00 2001
From: cubic <1198365636@qq.com>
Date: Sun, 13 Nov 2022 19:33:43 +0800
Subject: [PATCH] flatten.c correction related to output dims (#1381)

After flatten layer, the input dims 4 should be squeezed to 2. The original bug with output dims 4 is corrected.
---
 source/operator/prototype/flatten.c | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/operator/prototype/flatten.c b/source/operator/prototype/flatten.c
index 5354dfce..11536cc1 100644
--- a/source/operator/prototype/flatten.c
+++ b/source/operator/prototype/flatten.c
@@ -53,7 +53,9 @@ static int infer_shape(struct node* node)
 
     output->layout = TENGINE_LAYOUT_NHWC;
 
-    set_ir_tensor_shape(output, dims, 4);
+    //set_ir_tensor_shape(output, dims, 4);
+    //stephen@20220923
+    set_ir_tensor_shape(output, dims, 2);
 
     return 0;
 }
-- 
GitLab