From a4e138660270e7599793fa438cd7b2fc2ce215a6 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Thu, 29 Jul 2021 22:24:27 -0700 Subject: [PATCH] Add remaining validation to `sdca_internal.cc` PiperOrigin-RevId: 387738010 Change-Id: I28eedcfd87a53aaf34deb075acea1f8c95470808 --- tensorflow/core/kernels/sdca_internal.cc | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensorflow/core/kernels/sdca_internal.cc b/tensorflow/core/kernels/sdca_internal.cc index 6c4a63b270c25..164f9382724ca 100644 --- 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(); + 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(