Fix nullptr deref in tf.raw_ops.CTCLoss.

PiperOrigin-RevId: 372266334 Change-Id: Ic52c3e9f13a38f54482d670907eda1688450862b
tensorflow · May 6, 2021 · 14607c0 · 14607c0
1 parent 8410ce6
commit 14607c0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/ctc_loss_op.cc b/tensorflow/core/kernels/ctc_loss_op.cc
@@ -109,6 +109,9 @@ class CTCLossOp : public OpKernel {
 
     const TensorShape& inputs_shape = inputs->shape();
     const int64 max_time = inputs_shape.dim_size(0);
+    OP_REQUIRES(ctx, max_time != 0,
+                errors::InvalidArgument(
+                    "Max time or first dimension of input cannot be 0."));
     const int64 batch_size = inputs_shape.dim_size(1);
     const int64 num_classes_raw = inputs_shape.dim_size(2);
     OP_REQUIRES(