From 32fdcbff9d06d010d908fcc4bd4b36eb3ce15925 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Wed, 5 May 2021 22:39:29 -0700 Subject: [PATCH] Validate arguments of `FractionalMaxPoolGrad` PiperOrigin-RevId: 372274982 Change-Id: If46b0c442efa4eaef635ce6a476717060420122c --- tensorflow/core/kernels/fractional_max_pool_op.cc | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc index 619a3507ce415..1a2a783d135c5 100644 --- a/tensorflow/core/kernels/fractional_max_pool_op.cc +++ b/tensorflow/core/kernels/fractional_max_pool_op.cc @@ -235,6 +235,20 @@ class FractionalMaxPoolGradOp : public OpKernel { // Just to make it similar to FractionalMaxPoolOp. constexpr int tensor_in_and_out_dims = 4; + OP_REQUIRES( + context, tensor_in.dims() == tensor_in_and_out_dims, + errors::InvalidArgument("orig_input should be a tensor of rank 4, got ", + tensor_in.DebugString())); + OP_REQUIRES(context, tensor_in.NumElements() > 0, + errors::InvalidArgument("orig_input must not be empty, got ", + tensor_in.DebugString())); + OP_REQUIRES(context, tensor_out.dims() == tensor_in_and_out_dims, + errors::InvalidArgument( + "orig_output should be a tensor of rank 4, got ", + tensor_out.DebugString())); + OP_REQUIRES(context, tensor_out.NumElements() > 0, + errors::InvalidArgument("orig_output must not be empty, got ", + tensor_out.DebugString())); std::vector input_size(tensor_in_and_out_dims); std::vector output_size(tensor_in_and_out_dims); for (int i = 0; i < tensor_in_and_out_dims; ++i) {