diff options
Diffstat (limited to 'src/graph/nodes/FlattenLayer.cpp')
| -rw-r--r-- | src/graph/nodes/FlattenLayer.cpp | 6 |
1 files changed, 4 insertions, 2 deletions
diff --git a/src/graph/nodes/FlattenLayer.cpp b/src/graph/nodes/FlattenLayer.cpp index 66840f65f1..ea08296ba3 100644 --- a/src/graph/nodes/FlattenLayer.cpp +++ b/src/graph/nodes/FlattenLayer.cpp @@ -39,7 +39,9 @@ std::unique_ptr<arm_compute::IFunction> FlattenLayer::instantiate_node(GraphCont arm_compute::ITensor *out = output->tensor(); // Auto configure output - arm_compute::auto_init_if_empty(*out->info(), in->info()->tensor_shape(), 1, in->info()->data_type(), in->info()->fixed_point_position()); + TensorShape tensor_shape = in->info()->tensor_shape(); + tensor_shape.collapse(in->info()->num_dimensions()); + arm_compute::auto_init_if_empty(*out->info(), tensor_shape, 1, in->info()->data_type(), in->info()->fixed_point_position()); // Create node context NodeContext node_ctx(OperationType::FlattenLayer); @@ -49,4 +51,4 @@ std::unique_ptr<arm_compute::IFunction> FlattenLayer::instantiate_node(GraphCont // Get function return OperationRegistry::get().find_operation(OperationType::FlattenLayer, _target_hint)->configure(node_ctx); -} +}
\ No newline at end of file |