Merge pull request #51193 from tensorflow/mm-cherrypick-0f931751fb20f…

…565c4e94aa6df58d54a003cdb30-on-r2.4 Validate dimensions of input tensor in `FractionalAvgPoolGrad`
tensorflow · Aug 5, 2021 · c2533e1 · c2533e1
2 parents 101764b + 8d84da8
commit c2533e1
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/fractional_avg_pool_op.cc b/tensorflow/core/kernels/fractional_avg_pool_op.cc
@@ -271,6 +271,18 @@ class FractionalAvgPoolGradOp : public OpKernel {
     const int64 in_rows = orig_input_tensor_shape_flat(1);
     const int64 in_cols = orig_input_tensor_shape_flat(2);
     const int64 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