[tf.data] Use core::ScopedUnref to avoid resource leakage.

If for whatever reason iterator_resource->set_iterator did not return Status::OK(), we would leak a reference on the iterator_resource. With this change, we won't leak the resource.

PiperOrigin-RevId: 194662412

diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc
index f5db97f..a2f6c5f 100644 (file)
--- a/tensorflow/core/kernels/data/iterator_ops.cc
+++ b/tensorflow/core/kernels/data/iterator_ops.cc
@@ -584,9 +584,9 @@ class MakeIteratorOp : public OpKernel {
     IteratorResource* iterator_resource;
     OP_REQUIRES_OK(
         ctx, LookupResource(ctx, HandleFromInput(ctx, 1), &iterator_resource));
+    core::ScopedUnref unref(iterator_resource);
     OP_REQUIRES_OK(ctx, iterator_resource->set_iterator(
                             dataset->MakeIterator("Iterator")));
-    iterator_resource->Unref();
   }
 };