Fix tf.raw_ops.QuantizeAndDequantizeV4Grad CHECK failure.

PiperOrigin-RevId: 370532425 Change-Id: I767721be266851b63d8fe55e7ac6be0af6017f6c
tensorflow · Apr 26, 2021 · 20431e9 · 20431e9
1 parent 95078c1
commit 20431e9
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/quantize_and_dequantize_op.cc b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
@@ -160,7 +160,17 @@ class QuantizeAndDequantizeV4GradientOp : public OpKernel {
         errors::InvalidArgument("gradient and input must be the same size"));
     const int depth = (axis_ == -1) ? 1 : input.dim_size(axis_);
     const Tensor& input_min_tensor = ctx->input(2);
+    OP_REQUIRES(ctx,
+                input_min_tensor.dims() == 0 || input_min_tensor.dims() == 1,
+                errors::InvalidArgument(
+                    "Input min tensor must have dimension 1. Recieved ",
+                    input_min_tensor.dims(), "."));
     const Tensor& input_max_tensor = ctx->input(3);
+    OP_REQUIRES(ctx,
+                input_max_tensor.dims() == 0 || input_max_tensor.dims() == 1,
+                errors::InvalidArgument(
+                    "Input max tensor must have dimension 1. Recieved ",
+                    input_max_tensor.dims(), "."));
     if (axis_ != -1) {
       OP_REQUIRES(
           ctx, input_min_tensor.dim_size(0) == depth,