Merge pull request #49671 from geetachavan1/cherrypicks_WGAN4

Fix OOB read issue with `tf.raw_ops.CTCLoss`.
tensorflow · May 26, 2021 · ec4035f · ec4035f
2 parents a156675 + 67f1f7b
commit ec4035f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/ctc_loss_op.cc b/tensorflow/core/kernels/ctc_loss_op.cc
@@ -62,6 +62,10 @@ class CTCLossOp : public OpKernel {
                 errors::InvalidArgument("sequence_length is not a vector"));
     OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(labels_indices->shape()),
                 errors::InvalidArgument("labels_indices is not a matrix"));
+    OP_REQUIRES(ctx, labels_indices->dim_size(1) > 1,
+                errors::InvalidArgument(
+                    "labels_indices second dimension must be >= 1. Received ",
+                    labels_indices->dim_size(1)));
     OP_REQUIRES(ctx, TensorShapeUtils::IsVector(labels_values->shape()),
                 errors::InvalidArgument("labels_values is not a vector"));