From b22786e7e9b7bdb6a56936ff29cc7e9968d7bc1d Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Wed, 28 Apr 2021 15:31:26 -0700
Subject: [PATCH] Prevent division by 0

PiperOrigin-RevId: 370998952
Change-Id: I6b1d49079624ee1447d2d9b53a8976fb356cc8f5
---
 tensorflow/lite/kernels/split.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/lite/kernels/split.cc b/tensorflow/lite/kernels/split.cc
index be7425cc0b5e2..cf778b0b95124 100644
--- a/tensorflow/lite/kernels/split.cc
+++ b/tensorflow/lite/kernels/split.cc
@@ -60,6 +60,7 @@ TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node,
   TF_LITE_ENSURE(context, axis_value < NumDimensions(input));
 
   const int input_size = SizeOfDimension(input, axis_value);
+  TF_LITE_ENSURE(context, num_splits != 0);
   TF_LITE_ENSURE_MSG(context, input_size % num_splits == 0,
                      "Not an even split");
   const int slice_size = input_size / num_splits;