From c5b0d5f8ac19888e46ca14b0e27562e7fbbee9a9 Mon Sep 17 00:00:00 2001 From: Amit Patankar Date: Fri, 30 Apr 2021 10:39:05 -0700 Subject: [PATCH] Fix the CHECK failure in tf.raw_ops.QuantizeAndDequantizeV2. --- tensorflow/core/kernels/quantize_and_dequantize_op.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/core/kernels/quantize_and_dequantize_op.cc b/tensorflow/core/kernels/quantize_and_dequantize_op.cc index ae905e22..408196a2 100644 --- a/tensorflow/core/kernels/quantize_and_dequantize_op.cc +++ b/tensorflow/core/kernels/quantize_and_dequantize_op.cc @@ -72,6 +72,9 @@ class QuantizeAndDequantizeV2Op : public OpKernel { void Compute(OpKernelContext* ctx) override { const Tensor& input = ctx->input(0); + OP_REQUIRES( + ctx, axis_ >= -1, + errors::InvalidArgument("Axis must be at least -1. Found ", axis_)); OP_REQUIRES( ctx, (axis_ == -1 || axis_ < input.shape().dims()), errors::InvalidArgument("Shape must be at least rank", axis_ + 1, -- 2.27.0