Fix null ptr deref in tf.raw_ops.RaggedTensorToTensor

PiperOrigin-RevId: 384257511 Change-Id: I0484ad285039d132d6c41b284a7fcdd2b774a38e
tensorflow · Jul 12, 2021 · 301ae88 · 301ae88
1 parent e787d20
commit 301ae88
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
@@ -348,6 +348,9 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
   Status GetFirstDimensionSize(OpKernelContext* context, INDEX_TYPE* result) {
     const Tensor first_partition_tensor =
         context->input(kFirstPartitionInputIndex);
+    if (row_partition_types_.empty()) {
+      return errors::InvalidArgument("No row_partition_types given.");
+    }
     const RowPartitionType first_partition_type = row_partition_types_[0];
     switch (first_partition_type) {
       case RowPartitionType::FIRST_DIM_SIZE: