From e407b999409985d6b03fd70b0abe1e152b3460bb Mon Sep 17 00:00:00 2001 From: Jiri Simsa Date: Wed, 24 Feb 2021 13:45:00 -0800 Subject: [PATCH] [tf.data] Avoid accessing possibly uninitialized memory in the `RepeatDatasetOp` kernel. PiperOrigin-RevId: 359360786 Change-Id: Id83c789b5f24bac08ae13e1968dd1610fc29de01 --- .../core/kernels/data/repeat_dataset_op.cc | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tensorflow/core/kernels/data/repeat_dataset_op.cc b/tensorflow/core/kernels/data/repeat_dataset_op.cc index 4ed6bc63dca..06f268f9251 100644 --- a/tensorflow/core/kernels/data/repeat_dataset_op.cc +++ b/tensorflow/core/kernels/data/repeat_dataset_op.cc @@ -232,26 +232,26 @@ class RepeatDatasetOp::Dataset : public DatasetBase { TF_RETURN_IF_ERROR(dataset()->input_->MakeIterator( ctx, this, prefix(), &input_impl_)); } - Status s = input_impl_->GetNext(ctx, out_tensors, end_of_sequence); + TF_RETURN_IF_ERROR( + input_impl_->GetNext(ctx, out_tensors, end_of_sequence)); DCHECK(!*end_of_sequence || out_tensors->empty()); if (first_call_ && *end_of_sequence && !ctx->split_provider()) { - // If the first call to GetNext() fails because the end - // of sequence has been reached, we terminate the - // iteration immediately. (Otherwise, this iterator - // would loop infinitely and never produce a value.) + // If the first call to GetNext() fails because the end of sequence + // has been reached, we terminate the iteration immediately. + // Otherwise, this iterator would loop infinitely and never produce a + // value. input_impl_.reset(); return Status::OK(); } first_call_ = false; if (!*end_of_sequence) { - return s; - } else { - if (ctx->split_provider()) { - TF_RETURN_IF_ERROR(ctx->split_provider()->Reset()); - } - input_impl_.reset(); - first_call_ = true; + return Status::OK(); } + if (ctx->split_provider()) { + TF_RETURN_IF_ERROR(ctx->split_provider()->Reset()); + } + input_impl_.reset(); + first_call_ = true; } while (true); }