Fix SEGV in CTC ops

PiperOrigin-RevId: 372430279 Change-Id: I7ec2ad9d6f4d0980c33de45d27c6b17df5c6e26f
tensorflow · May 6, 2021 · b1b3230 · b1b3230
1 parent a74768f
commit b1b3230
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/ctc_decoder_ops.cc b/tensorflow/core/kernels/ctc_decoder_ops.cc
@@ -70,6 +70,9 @@ class CTCDecodeHelper {
     if (inputs_shape.dims() != 3) {
       return errors::InvalidArgument("inputs is not a 3-Tensor");
     }
+    if (inputs_shape.num_elements() == 0) {
+      return errors::InvalidArgument("inputs must not be empty");
+    }
 
     const int64 max_time = inputs_shape.dim_size(0);
     const int64 batch_size = inputs_shape.dim_size(1);