From 106d8f4fb89335a2c52d7c895b7a7485465ca8d9 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Tue, 27 Apr 2021 17:47:36 -0700 Subject: [PATCH] Prevent division by 0 in TFLite --- tensorflow/lite/kernels/depth_to_space.cc | 1 + tensorflow/lite/kernels/depth_to_space_test.cc | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/tensorflow/lite/kernels/depth_to_space.cc b/tensorflow/lite/kernels/depth_to_space.cc index 1637ad43..c2047f10 100644 --- a/tensorflow/lite/kernels/depth_to_space.cc +++ b/tensorflow/lite/kernels/depth_to_space.cc @@ -58,6 +58,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type); const int block_size = params->block_size; + TF_LITE_ENSURE(context, block_size > 0); const int input_height = input->dims->data[1]; const int input_width = input->dims->data[2]; const int input_channels = input->dims->data[3]; diff --git a/tensorflow/lite/kernels/depth_to_space_test.cc b/tensorflow/lite/kernels/depth_to_space_test.cc index 4429faf9..e0de01a9 100644 --- a/tensorflow/lite/kernels/depth_to_space_test.cc +++ b/tensorflow/lite/kernels/depth_to_space_test.cc @@ -60,6 +60,11 @@ TEST(DepthToSpaceOpModel, BadBlockSize) { EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 4), "Cannot allocate tensors"); } + +TEST(DepthToSpaceOpModel, NoBlockSize) { + EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 0), + "Cannot allocate tensors"); +} #endif TEST(DepthToSpaceOpModel, Float32) { -- 2.23.0