Skip to content
Permalink
Browse files Browse the repository at this point in the history
Validate dimensions of input tensor in FractionalAvgPoolGrad
PiperOrigin-RevId: 388286227
Change-Id: Ieb7566155e92acc8993a2212c76deacadc0edc8a
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Aug 2, 2021
1 parent 55e763f commit 0f93175
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions tensorflow/core/kernels/fractional_avg_pool_op.cc
Expand Up @@ -271,6 +271,18 @@ class FractionalAvgPoolGradOp : public OpKernel {
const int64_t in_rows = orig_input_tensor_shape_flat(1);
const int64_t in_cols = orig_input_tensor_shape_flat(2);
const int64_t in_depth = orig_input_tensor_shape_flat(3);
OP_REQUIRES(
context, in_batch != 0,
errors::InvalidArgument("Batch dimension of input must not be 0"));
OP_REQUIRES(
context, in_rows != 0,
errors::InvalidArgument("Rows dimension of input must not be 0"));
OP_REQUIRES(
context, in_cols != 0,
errors::InvalidArgument("Columns dimension of input must not be 0"));
OP_REQUIRES(
context, in_depth != 0,
errors::InvalidArgument("Depth dimension of input must not be 0"));

constexpr int tensor_in_and_out_dims = 4;
// Transform orig_input_tensor_shape into TensorShape
Expand Down

0 comments on commit 0f93175

Please sign in to comment.