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

[tf.data] Solve the bug when output status not OK, for PrefetchDataset, RecordBufferEnqueue function may be called more times than RecordBufferDequeue.

Browse Source 
PiperOrigin-RevId: 338149161
Change-Id: I8a7df4e21364814c6097447465d26f145219b776
This commit is contained in:
Jay Shi 2020-10-20 15:22:39 -07:00 committed by TensorFlower Gardener
parent 776e574e2a
commit 833b3a49a9
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/core/kernels/data/prefetch_dataset_op.cc
View File

@ -369,6 +369,11 @@ class PrefetchDatasetOp::Dataset : public DatasetBase {
}
*out_tensors = std::move(buffer_.front().value);
RecordBufferDequeue(ctx, *out_tensors);
} else {
// If status not ok, we still record the dequeue event to make sure each
// enqueue event is paired with a dequeue event even in the presence of
// errors.
RecordBufferDequeue(ctx, buffer_.front().value);
}
if (legacy_autotune_) {
auto_tuner_.RecordConsumption(buffer_.size());