diff options
Diffstat (limited to 'tensorflow/core/kernels/function_ops.cc')
| -rw-r--r-- | tensorflow/core/kernels/function_ops.cc | 12 |
1 files changed, 12 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc index e269aa3a78..d5c33c0188 100644 --- a/tensorflow/core/kernels/function_ops.cc +++ b/tensorflow/core/kernels/function_ops.cc @@ -135,6 +135,12 @@ REGISTER_KERNEL_BUILDER(Name(kArgOp) .TypeConstraint<ResourceHandle>("T"), ArgOp); +REGISTER_KERNEL_BUILDER(Name(kArgOp) + .Device(DEVICE_GPU) + .HostMemory("output") + .TypeConstraint<string>("T"), + ArgOp); + #define REGISTER(type) \ REGISTER_KERNEL_BUILDER( \ Name(kRetOp).Device(DEVICE_GPU).TypeConstraint<type>("T"), RetvalOp); @@ -149,6 +155,12 @@ REGISTER_KERNEL_BUILDER(Name(kRetOp) .TypeConstraint<ResourceHandle>("T") .HostMemory("input"), RetvalOp); + +REGISTER_KERNEL_BUILDER(Name(kRetOp) + .Device(DEVICE_GPU) + .TypeConstraint<string>("T") + .HostMemory("input"), + RetvalOp); #undef REGISTER class PassOn : public OpKernel { |