diff options
| author |  Brennan Saeta <saeta@google.com> | 2018-04-28 10:51:32 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-04-28 10:54:57 -0700 |
| commit | fb1069781ffcbac222392a68c01a45fae264888e (patch) | |
| tree | 45c55f7d832be4f6b278787424c423c242f15d4f | |
| parent | c01858350a1fc0f0fbf9a38fcd5c71e565343316 (diff) | |
[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
| -rw-r--r-- | tensorflow/core/kernels/data/iterator_ops.cc | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc index f5db97fd59..a2f6c5fe2c 100644 --- 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(); } }; |