From 2c74674348a4708ced58ad6eb1b23354df8ee044 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Wed, 28 Apr 2021 13:57:37 -0700 Subject: [PATCH] Prevent division by 0 PiperOrigin-RevId: 370979352 Change-Id: Ic79191c316d986fc6072ecaebfec9d5f2b924d00 --- tensorflow/lite/kernels/batch_to_space_nd.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/lite/kernels/batch_to_space_nd.cc b/tensorflow/lite/kernels/batch_to_space_nd.cc index 9d6492e0fcbf0..044ac1b3a5ee5 100644 --- a/tensorflow/lite/kernels/batch_to_space_nd.cc +++ b/tensorflow/lite/kernels/batch_to_space_nd.cc @@ -78,6 +78,7 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context, int output_batch_size = input_size->data[0]; for (int dim = 0; dim < spatial_dims_num; ++dim) { // Number of batch must be multiple of (block_shape[dim]). + TF_LITE_ENSURE(context, block_shape[dim] != 0); TF_LITE_ENSURE_EQ(context, output_batch_size % block_shape[dim], 0); output_batch_size = output_batch_size / block_shape[dim]; output_size->data[dim + 1] = input_size->data[dim + 1] * block_shape[dim] -