From 2c74674348a4708ced58ad6eb1b23354df8ee044 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
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] -