diff options
| -rw-r--r-- | tensorflow/core/kernels/resource_variable_ops.cc | 18 |
1 files changed, 11 insertions, 7 deletions
diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc index 427044ca67..23d76986bf 100644 --- a/tensorflow/core/kernels/resource_variable_ops.cc +++ b/tensorflow/core/kernels/resource_variable_ops.cc @@ -172,17 +172,21 @@ REGISTER_KERNEL_BUILDER( .Device(DEVICE_GPU) \ .HostMemory("resource") \ .TypeConstraint<type>("dtype"), \ - ResourceHandleOp<Var>) \ - REGISTER_KERNEL_BUILDER(Name("_VarHandlesOp") \ - .Device(DEVICE_GPU) \ - .HostMemory("resources") \ - .TypeConstraint<type>("dtypes"), \ - ResourceHandlesOp<Var>) - + ResourceHandleOp<Var>) TF_CALL_GPU_ALL_TYPES(REGISTER_GPU_KERNELS); TF_CALL_int64(REGISTER_GPU_KERNELS); TF_CALL_variant(REGISTER_GPU_KERNELS); #undef REGISTER_GPU_KERNELS + +REGISTER_KERNEL_BUILDER(Name("_VarHandlesOp") + .Device(DEVICE_GPU) + .HostMemory("resources") + .TypeConstraint("dtypes", + {DT_INT64, DT_COMPLEX64, + DT_COMPLEX128, DT_HALF, DT_FLOAT, + DT_DOUBLE, DT_BOOL, DT_VARIANT}), + ResourceHandlesOp<Var>); + #endif // GOOGLE_CUDA template <typename T> |