From bf686faeddcca97be6ad7b6421cb26ab1c3cea2c Mon Sep 17 00:00:00 2001 From: Yu-Cheng Ling Date: Wed, 20 Mar 2019 11:35:38 -0700 Subject: [PATCH] TFLite: Enhance input check for ResizeNearestNeghbor PiperOrigin-RevId: 239440642 --- tensorflow/lite/kernels/resize_nearest_neighbor.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/kernels/resize_nearest_neighbor.cc b/tensorflow/lite/kernels/resize_nearest_neighbor.cc index 3030a4f28e2..16db0997806 100644 --- a/tensorflow/lite/kernels/resize_nearest_neighbor.cc +++ b/tensorflow/lite/kernels/resize_nearest_neighbor.cc @@ -57,11 +57,13 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { const TfLiteTensor* size = GetInput(context, node, kSizeTensor); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - // TODO(ahentz): Our current implementations rely on the inputs being 4D. + // TODO(ahentz): Our current implementations rely on the input being 4D, + // and the size being 1D tensor with exactly 2 elements. TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4); TF_LITE_ENSURE_EQ(context, NumDimensions(size), 1); - TF_LITE_ENSURE_EQ(context, size->type, kTfLiteInt32); + TF_LITE_ENSURE_EQ(context, size->dims->data[0], 2); + output->type = input->type; if (!IsConstantTensor(size)) {