Adjust checks in reshape (#2383)

BUG=https://github.com/tensorflow/tflite-micro/issues/2382

tensorflow/lite/micro/kernels/reshape_common.cc

diff --git a/tensorflow/lite/micro/kernels/reshape_common.cc b/tensorflow/lite/micro/kernels/reshape_common.cc
index 91fd1bc..75bc2eb 100644
--- a/tensorflow/lite/micro/kernels/reshape_common.cc
+++ b/tensorflow/lite/micro/kernels/reshape_common.cc
@@ -1,4 +1,4 @@
-/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -52,9 +52,9 @@
   if (new_shape != nullptr && new_shape->dims->size > 0) {
     // use new shape tensor data
     TF_LITE_ENSURE_EQ(context, new_shape->type, kTfLiteInt32);
-    TF_LITE_ENSURE_EQ(context, new_shape->dims->size, 1);
     output_shape_data = GetTensorData<int>(new_shape);
-    output_shape_size = new_shape->dims->data[0];
+    output_shape_size = new_shape->dims->data[new_shape->dims->size - 1];
+
     TF_LITE_ENSURE_EQ(context, output_shape_size,
                       static_cast<int32_t>(output->dims->size));
   } else {