From b22786e7e9b7bdb6a56936ff29cc7e9968d7bc1d Mon Sep 17 00:00:00 2001 From: Mihai Maruseac 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;