experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

TFLite: Enhance input check for ResizeNearestNeghbor

Browse Source 
PiperOrigin-RevId: 239440642
This commit is contained in:
Yu-Cheng Ling 2019-03-20 11:35:38 -07:00 committed by TensorFlower Gardener
parent 41dac366fb
commit bf686faedd
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/lite/kernels/resize_nearest_neighbor.cc
View File

@ -57,11 +57,13 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* size = GetInput(context, node, kSizeTensor); const TfLiteTensor* size = GetInput(context, node, kSizeTensor);
TfLiteTensor* output = GetOutput(context, node, kOutputTensor); 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(input), 4);
TF_LITE_ENSURE_EQ(context, NumDimensions(size), 1); TF_LITE_ENSURE_EQ(context, NumDimensions(size), 1);
TF_LITE_ENSURE_EQ(context, size->type, kTfLiteInt32); TF_LITE_ENSURE_EQ(context, size->type, kTfLiteInt32);
TF_LITE_ENSURE_EQ(context, size->dims->data[0], 2);
output->type = input->type; output->type = input->type;
if (!IsConstantTensor(size)) { if (!IsConstantTensor(size)) {