Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent division by 0.
PiperOrigin-RevId: 370962554
Change-Id: I0b9b62f4d8e1046dd88f9433f8dfeaf61a901680
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Apr 28, 2021
1 parent 09c73bc commit ff489d9
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions tensorflow/lite/kernels/conv.cc
Expand Up @@ -545,6 +545,7 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
// Only one scale factor per batch is typically necessary. See optimized
// implementation for why we need to allocate for the height of the inputs
// flattened to 2D.
TF_LITE_ENSURE(context, channels_in != 0);
const int height = NumElements(input) / channels_in;
int scaling_dims[1] = {height};
if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) {
Expand Down Expand Up @@ -587,6 +588,7 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
input_offsets->type = kTfLiteInt32;
input_offsets->allocation_type = kTfLiteArenaRw;
// See above comment for the need to allocate for height of inputs.
TF_LITE_ENSURE(context, channels_in != 0);
const int height = NumElements(input) / channels_in;
const int input_offset_dims[1] = {height};
if (!TfLiteIntArrayEqualsArray(input_offsets->dims, 1,
Expand Down Expand Up @@ -886,8 +888,9 @@ TfLiteStatus EvalHybridPerChannel(TfLiteContext* context, TfLiteNode* node,
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;
TfLiteTensor* quantized_input_tensor;
TF_LITE_ENSURE_OK(context,
GetTemporarySafe(context, node, data->input_quantized_index,
Expand Down Expand Up @@ -989,8 +992,9 @@ TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node,
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 float* input_ptr = GetTensorData<float>(input);
TfLiteTensor* quantized_input_tensor;
Expand Down

0 comments on commit ff489d9

Please sign in to comment.