diff options
| author |  Igor Ganichev <iga@google.com> | 2018-07-12 11:06:21 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-12 11:09:50 -0700 |
| commit | d843a4eee58314320354cbd53e80a53a78449da6 (patch) | |
| tree | 06e43e1a1349e87824edd8c0d62301884abf9aab /tensorflow/core/kernels/resource_variable_ops.h | |
| parent | f8ce8dc04f569b3bbd1ac7cc9e358a651530990d (diff) | |
Register DestroyResourceOp for XLA devices
Before this change, we were not releasing device memory
allocated by ResourceVariables.
PiperOrigin-RevId: 204329027
Diffstat (limited to 'tensorflow/core/kernels/resource_variable_ops.h')
| -rw-r--r-- | tensorflow/core/kernels/resource_variable_ops.h | 9 |
1 files changed, 9 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/resource_variable_ops.h b/tensorflow/core/kernels/resource_variable_ops.h index 8cae5d21f0..9b60106f13 100644 --- a/tensorflow/core/kernels/resource_variable_ops.h +++ b/tensorflow/core/kernels/resource_variable_ops.h @@ -28,6 +28,15 @@ class ReadVariableOp : public OpKernel { DataType dtype_; }; +class DestroyResourceOp : public OpKernel { + public: + explicit DestroyResourceOp(OpKernelConstruction* ctx); + void Compute(OpKernelContext* ctx) override; + + private: + bool ignore_lookup_error_; +}; + } // namespace tensorflow #endif // TENSORFLOW_CORE_KERNELS_RESOURCE_VARIABLE_OPS_H_ |