From 6d36ba65577006affb272335b7c1abd829010708 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Wed, 28 Apr 2021 14:22:30 -0700
Subject: [PATCH] Prevent division by 0

PiperOrigin-RevId: 370984990
Change-Id: Ib324955bbeb1cbd97c82fd5d61a00a2697c9a2de
---
 tensorflow/lite/kernels/space_to_batch_nd.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/lite/kernels/space_to_batch_nd.cc b/tensorflow/lite/kernels/space_to_batch_nd.cc
index 0d537e2d1892f..af7b9d9e914a1 100644
--- a/tensorflow/lite/kernels/space_to_batch_nd.cc
+++ b/tensorflow/lite/kernels/space_to_batch_nd.cc
@@ -79,6 +79,7 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
   for (int dim = 0; dim < spatial_dims_num; ++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];
     output_batch_size *= block_shape[dim];