COMPMID-556: Fix checks in NEGEMMLowpOffsetContributionKernel

Change-Id: I767b7c941860cb372d7f42f6b747cb610abc223d
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/112580
Tested-by: BSG Visual Compute Jenkins server to access repositories on http://mpd-gerrit.cambridge.arm.com <bsgcomp@arm.com>
Reviewed-by: Gian Marco Iodice <gianmarco.iodice@arm.com>

1 files changed, 15 insertions, 12 deletions

diff --git a/src/core/NEON/kernels/NEGEMMLowpOffsetContributionKernel.cpp b/src/core/NEON/kernels/NEGEMMLowpOffsetContributionKernel.cpp
index 2cf99dd292..f6964002dd 100644
--- a/src/core/NEON/kernels/NEGEMMLowpOffsetContributionKernel.cpp
+++ b/src/core/NEON/kernels/NEGEMMLowpOffsetContributionKernel.cpp
@@ -64,21 +64,24 @@ Status validate_arguments(const ITensorInfo *mm_result, const ITensorInfo *vecto
         ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(vector_sum_row, 1, DataType::S32);
         ARM_COMPUTE_RETURN_ERROR_ON(vector_sum_row->dimension(0) != mm_result->dimension(1));
 
-        TensorShape output_shape         = mm_result->tensor_shape();
-        TensorShape vector_sum_row_shape = vector_sum_row->tensor_shape();
-        vector_sum_row_shape.collapse(1);
-        output_shape.collapse(2);
+        TensorShape output_shape = mm_result->tensor_shape();
+        if(output_shape.num_dimensions() > 1)
+        {
+            TensorShape vector_sum_row_shape = vector_sum_row->tensor_shape();
+            vector_sum_row_shape.collapse_from(1);
+            output_shape.collapse_from(2);
 
-        ARM_COMPUTE_RETURN_ERROR_ON_MSG(vector_sum_row_shape[1] != output_shape[2],
-                                        "mm_result tensor must have the same number of batches of output tensor");
+            ARM_COMPUTE_ERROR_ON_MSG(vector_sum_row_shape[1] != output_shape[2], "mm_result tensor must have the same number of batches of output tensor");
 
-        if(a_offset != 0)
-        {
-            TensorShape vector_sum_col_shape = vector_sum_col->tensor_shape();
-            vector_sum_col_shape.collapse(1);
+            if(a_offset != 0)
+            {
+                TensorShape vector_sum_col_shape = vector_sum_col->tensor_shape();
+                vector_sum_col_shape.collapse_from(1);
 
-            ARM_COMPUTE_RETURN_ERROR_ON_MSG(vector_sum_col_shape[1] != 1 && vector_sum_col_shape[1] != vector_sum_row_shape[1],
-                                            "vector_sum_col tensor must have the same number of batches of vector_sum_row_shape or the number of batches must be set to 1");
+                ARM_COMPUTE_ERROR_ON_MSG(vector_sum_col_shape[1] != 1
+                                         && vector_sum_col_shape[1] != vector_sum_row_shape[1],
+                                         "vector_sum_col tensor must have the same number of batches of vector_sum_row_shape or the number of batches must be set to 1");
+            }
         }
     }