Merge pull request #52801 from pranve/cherrypick-25e731a020a1768ce8bea952180b9e35d081c6ff-on-r2.6

mihaimaruseac · web-flow · commit a35dad149588 · 2021-10-28T15:02:27.000-07:00
[tf.data] Fix memory leak in to_tf_record_op.
diff --git a/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc b/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc
@@ -87,6 +87,7 @@ class ToTFRecordOp : public AsyncOpKernel {
     IteratorContext iter_ctx(std::move(params));
     DatasetBase* finalized_dataset;
     TF_RETURN_IF_ERROR(FinalizeDataset(ctx, dataset, &finalized_dataset));
+    core::ScopedUnref unref(finalized_dataset);
 
     std::unique_ptr<IteratorBase> iterator;
     TF_RETURN_IF_ERROR(finalized_dataset->MakeIterator(
