Merge pull request #49781 from geetachavan1/cherrypicks_V4G18

Prevent division by 0
tensorflow · May 30, 2021 · 0792334 · 0792334
2 parents d41015b + 3e1f82c
commit 0792334
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tensorflow/lite/kernels/space_to_batch_nd.cc b/tensorflow/lite/kernels/space_to_batch_nd.cc
@@ -75,6 +75,7 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
   for (int dim = 0; dim < kSpatialDimensionNum; ++dim) {
     int final_dim_size = (input_size->data[dim + 1] + paddings_data[dim * 2] +
                           paddings_data[dim * 2 + 1]);
+    TF_LITE_ENSURE(context, block_shape[dim] != 0);
     TF_LITE_ENSURE_EQ(context, final_dim_size % block_shape[dim], 0);
     output_size->data[dim + 1] = final_dim_size / block_shape[dim];
   }