diff options
| -rw-r--r-- | tensorflow/core/kernels/resource_variable_ops.cc | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc index 3cca493972..90db0c2b7b 100644 --- a/tensorflow/core/kernels/resource_variable_ops.cc +++ b/tensorflow/core/kernels/resource_variable_ops.cc @@ -200,6 +200,9 @@ class DestroyResourceOp : public OpKernel { REGISTER_KERNEL_BUILDER(Name("DestroyResourceOp").Device(DEVICE_CPU), DestroyResourceOp); +REGISTER_KERNEL_BUILDER( + Name("DestroyResourceOp").Device(DEVICE_GPU).HostMemory("resource"), + DestroyResourceOp); template <typename Device, typename T> class AssignVariableOp : public OpKernel { |