diff options
| author |  Peter Hawkins <phawkins@google.com> | 2018-09-28 17:11:09 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-28 17:15:27 -0700 |
| commit | 05e5d2a1a9d5471d634043135834ecae4355926a (patch) | |
| tree | 962d0f088a13c38ffc5fd47bf66036ca025c51fa /tensorflow/core/kernels | |
| parent | 6d354f6bd686d748d02039f26197f590b817b8c3 (diff) | |
[TF] Fix incorrect type constraint on _VarHandlesOp kernel on GPU. The kernel supports any combination of the dtypes, and does not need a separate kernel for each dtype.
PiperOrigin-RevId: 215019812
Diffstat (limited to 'tensorflow/core/kernels')
| -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> |