Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix out of bound access in DequantizeOp by adding check for axis < in…
…put dimension

PiperOrigin-RevId: 411214268
Change-Id: I3249d2a69ddc82f182c589a3a5bbfb71543f4b29
  • Loading branch information
ishark authored and tensorflower-gardener committed Nov 20, 2021
1 parent 5665767 commit 23968a8
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions tensorflow/core/kernels/dequantize_op.cc
Expand Up @@ -94,6 +94,11 @@ class DequantizeOp : public OpKernel {
const Tensor& input_min_tensor = ctx->input(1);
const Tensor& input_max_tensor = ctx->input(2);

OP_REQUIRES(
ctx, axis_ < input.dims(),
errors::InvalidArgument("Axis must be less than input dimension(",
input.dims(), "), got ", axis_));

int num_slices = 1;
if (axis_ > -1) {
num_slices = input.dim_size(axis_);
Expand Down

0 comments on commit 23968a8

Please sign in to comment.