From cbda3c6b2dbbd3fbdc482ff8c0170a78ec2e97d0 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Wed, 28 Apr 2021 15:53:48 -0700 Subject: [PATCH] Prevent divisions by 0 --- tensorflow/lite/kernels/depthwise_conv.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/kernels/depthwise_conv.cc b/tensorflow/lite/kernels/depthwise_conv.cc index 3f75287e..fb031926 100644 --- a/tensorflow/lite/kernels/depthwise_conv.cc +++ b/tensorflow/lite/kernels/depthwise_conv.cc @@ -276,6 +276,7 @@ TfLiteStatus ComputeDepthMultiplier(TfLiteContext* context, int16* depth_multiplier) { int num_filter_channels = SizeOfDimension(filter, 3); int num_input_channels = SizeOfDimension(input, 3); + TF_LITE_ENSURE(context, num_input_channels != 0); TF_LITE_ENSURE_EQ(context, num_filter_channels % num_input_channels, 0); *depth_multiplier = num_filter_channels / num_input_channels; @@ -446,8 +447,9 @@ TfLiteStatus EvalHybridPerChannel(TfLiteContext* context, TfLiteNode* node, float output_activation_min, output_activation_max; CalculateActivationRange(params->activation, &output_activation_min, &output_activation_max); - const int input_size = NumElements(input) / SizeOfDimension(input, 0); const int batch_size = SizeOfDimension(input, 0); + TF_LITE_ENSURE(context, batch_size != 0); + const int input_size = NumElements(input) / batch_size; const TfLiteTensor* input_quantized = GetTemporary(context, node, data->input_quantized_index); int8_t* quantized_input_ptr_batch = input_quantized->data.int8; -- 2.23.0