diff options
author | 2018-09-24 19:26:46 -0700 | |
---|---|---|
committer | 2018-09-24 19:31:03 -0700 | |
commit | bb1c131aad55e336d25fd297ecd8582773d6476f (patch) | |
tree | 46daf6e7ae9bc3945947b4cfcdff2e3fa5b67c01 /tensorflow/core/kernels | |
parent | 720594142a51c2676d086ba00705d95002474687 (diff) |
Fix memory leak of a Var resource in the multiple variable-handling kernels.
This change fixes memory leaks in the ScatterNdUpdateOp and StridedSliceAssign kernels, and in training-op kernels that use `GetTrainingVariableMutex()`.
PiperOrigin-RevId: 214372346
Diffstat (limited to 'tensorflow/core/kernels')
-rw-r--r-- | tensorflow/core/kernels/scatter_nd_op.cc | 1 | ||||
-rw-r--r-- | tensorflow/core/kernels/strided_slice_op.cc | 1 | ||||
-rw-r--r-- | tensorflow/core/kernels/training_op_helpers.cc | 1 |
3 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/scatter_nd_op.cc b/tensorflow/core/kernels/scatter_nd_op.cc index e0194605ce..2f8aede427 100644 --- a/tensorflow/core/kernels/scatter_nd_op.cc +++ b/tensorflow/core/kernels/scatter_nd_op.cc @@ -145,6 +145,7 @@ class ScatterNdUpdateOp : public OpKernel { if (dtype_ == DT_RESOURCE) { Var* v; OP_REQUIRES_OK(c, LookupResource(c, HandleFromInput(c, 0), &v)); + core::ScopedUnref scoped_unref(v); mutex_lock m(*v->mu()); DoCompute(c); } else if (use_exclusive_lock_) { diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc index 7b537fef5b..f0575de4d9 100644 --- a/tensorflow/core/kernels/strided_slice_op.cc +++ b/tensorflow/core/kernels/strided_slice_op.cc @@ -306,6 +306,7 @@ class StridedSliceAssignOp : public OpKernel { Var* v; OP_REQUIRES_OK(context, LookupResource(context, HandleFromInput(context, 0), &v)); + core::ScopedUnref scoped_unref(v); mutex_lock ml(*v->mu()); OP_REQUIRES_OK(context, PrepareToUpdateVariable<Device, T>(context, v->tensor())); diff --git a/tensorflow/core/kernels/training_op_helpers.cc b/tensorflow/core/kernels/training_op_helpers.cc index d3c4f62071..83b83fcdb9 100644 --- a/tensorflow/core/kernels/training_op_helpers.cc +++ b/tensorflow/core/kernels/training_op_helpers.cc @@ -21,6 +21,7 @@ mutex* GetTrainingVariableMutex(OpKernelContext* ctx, int input) { if (ctx->input_dtype(input) == DT_RESOURCE) { Var* var; if (LookupResource(ctx, HandleFromInput(ctx, input), &var).ok()) { + core::ScopedUnref scoped_unref(var); return var->mu(); } else { ctx->CtxFailureWithWarning( |