diff options
| author |  Nupur Garg <nupurgarg@google.com> | 2018-01-30 15:11:33 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-01-30 15:19:19 -0800 |
| commit | e7df91c5f65e8b10535d1a09c817aa4e900e4001 (patch) | |
| tree | 2e3bab64177efcc8ac566cc20afac2904bf88d22 /tensorflow/contrib/lite/kernels/batch_to_space_nd.cc | |
| parent | ac7167794f7c22d42c14a8d2f47b9533e118d73f (diff) | |
Internal change.
PiperOrigin-RevId: 183900332
Diffstat (limited to 'tensorflow/contrib/lite/kernels/batch_to_space_nd.cc')
| -rw-r--r-- | tensorflow/contrib/lite/kernels/batch_to_space_nd.cc | 9 |
1 files changed, 8 insertions, 1 deletions
diff --git a/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc b/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc index d84a77039b..889239f932 100644 --- a/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc +++ b/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc @@ -57,6 +57,7 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context, BatchToSpaceNDContext* op_context) { TfLiteIntArray* input_size = op_context->input->dims; const int* block_shape = GetTensorData<int32>(op_context->block_shape); + const int* crops = GetTensorData<int32>(op_context->crops); TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->block_shape), kBlockSizeDimensionNum); @@ -65,7 +66,13 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context, TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->crops), kSpatialDimensionNum); - // TODO(ycling): Add crops as part of calculation. + // TODO(ycling): Add crops as part of calculation. Remove check for a crops + // containing all zeroes. + TF_LITE_ENSURE_EQ(context, crops[0], 0); + TF_LITE_ENSURE_EQ(context, crops[1], 0); + TF_LITE_ENSURE_EQ(context, crops[2], 0); + TF_LITE_ENSURE_EQ(context, crops[3], 0); + // Number of batch must be multiple of (block_shape[0] * block_shape[1]). TF_LITE_ENSURE_EQ(context, input_size->data[0] % (block_shape[0] * block_shape[1]), 0); |