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

[tf.data] Change the code to solve the bug that ParallelMap records more buffered elements enqueue than dequeue.

Browse Source 
PiperOrigin-RevId: 335742261
Change-Id: I5f43c880f19022b59addc02dd9a93da74ad96a74
This commit is contained in:
Jay Shi 2020-10-06 16:12:43 -07:00 committed by TensorFlower Gardener
parent 27134a78b3
commit bde6a9cfee
1 changed files with 1 additions and 1 deletions

2
tensorflow/core/kernels/data/parallel_map_dataset_op.cc
View File

@ -405,7 +405,6 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
TF_LOCKS_EXCLUDED(*mu_) {
mutex_lock l(*mu_);
num_calls_--;
RecordBufferEnqueue(ctx.get(), result->return_values);
result->notification.Notify();
cond_var_->notify_all();
}
@ -428,6 +427,7 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
auto done = [this, ctx, result](Status status) {
result->status.Update(status);
RecordBufferEnqueue(ctx.get(), result->return_values);
CallCompleted(ctx, result);
};