diff options
| author |  Alexandre Passos <apassos@google.com> | 2018-02-15 12:39:52 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-02-15 12:44:26 -0800 |
| commit | 4b297b5434438175b016da05421e7ddd46c0f8ee (patch) | |
| tree | f254da3f2c54b277ee676b8121506daa3a76f86a /tensorflow/core/kernels/dense_update_ops.cc | |
| parent | 4ed183d9d471ca04cf3961610a027136298c1788 (diff) | |
Register kernels for Assign and AssignVariableOp on GPU for integer types.
PiperOrigin-RevId: 185882834
Diffstat (limited to 'tensorflow/core/kernels/dense_update_ops.cc')
| -rw-r--r-- | tensorflow/core/kernels/dense_update_ops.cc | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/dense_update_ops.cc b/tensorflow/core/kernels/dense_update_ops.cc index 6497c8f371..0de97de205 100644 --- a/tensorflow/core/kernels/dense_update_ops.cc +++ b/tensorflow/core/kernels/dense_update_ops.cc @@ -109,6 +109,7 @@ TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS); AssignOpT<GPUDevice, type>); TF_CALL_GPU_ALL_TYPES(REGISTER_GPU_KERNELS); +TF_CALL_int64(REGISTER_GPU_KERNELS); #undef REGISTER_GPU_KERNELS #endif // GOOGLE_CUDA @@ -142,6 +143,7 @@ TF_CALL_NUMBER_TYPES(REGISTER_KERNELS); Name("AssignSub").Device(DEVICE_GPU).TypeConstraint<type>("T"), \ DenseUpdateOp<GPUDevice, type, DenseUpdateType::SUB>); TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS); +TF_CALL_int64(REGISTER_GPU_KERNELS); #undef REGISTER_GPU_KERNELS #endif // end GOOGLE_CUDA |