From bde6a9cfeee634889e8e844cf8e30cf012a734ae Mon Sep 17 00:00:00 2001
From: Jay Shi <xiaojies@google.com>
Date: Tue, 6 Oct 2020 16:12:43 -0700
Subject: [PATCH] [tf.data] Change the code to solve the bug that ParallelMap
 records more buffered elements enqueue than dequeue.

PiperOrigin-RevId: 335742261
Change-Id: I5f43c880f19022b59addc02dd9a93da74ad96a74
---
 tensorflow/core/kernels/data/parallel_map_dataset_op.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
index 87ea4531d5d..8c1a49d9dda 100644
--- a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
@@ -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);
       };