experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[tf.data] Avoid accessing possibly uninitialized memory in the RepeatDatasetOp kernel.

Browse Source 
PiperOrigin-RevId: 359360786
Change-Id: Id83c789b5f24bac08ae13e1968dd1610fc29de01
This commit is contained in:
Jiri Simsa 2021-02-24 13:45:00 -08:00 committed by TensorFlower Gardener
parent 5a7974d25f
commit e407b99940
1 changed files with 12 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
tensorflow/core/kernels/data/repeat_dataset_op.cc
View File

@ -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);
}