Add remaining validation to sdca_internal.cc

PiperOrigin-RevId: 387738010 Change-Id: I28eedcfd87a53aaf34deb075acea1f8c95470808
tensorflow · Jul 30, 2021 · a4e1386 · a4e1386
1 parent 578e634
commit a4e1386
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/sdca_internal.cc b/tensorflow/core/kernels/sdca_internal.cc
@@ -380,6 +380,11 @@ Status Examples::Initialize(OpKernelContext* const context,
   const Tensor* example_labels_t;
   TF_RETURN_IF_ERROR(context->input("example_labels", &example_labels_t));
   auto example_labels = example_labels_t->flat<float>();
+  if (example_labels.size() != num_examples) {
+    return errors::InvalidArgument("Expected ", num_examples,
+                                   " example labels but got ",
+                                   example_labels.size());
+  }
 
   OpInputList dense_features_inputs;
   TF_RETURN_IF_ERROR(