diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-08-31 07:30:57 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-31 07:37:17 -0700 |
| commit | 371a8dd2e1e18fd35f07ac230c52a471d90d3538 (patch) | |
| tree | 1476a762285886d1f2a688ed8c54df14d3f59e15 | |
| parent | d09bd33fc9175c7fdbdded8f98b5c5d3a9f8ad7d (diff) | |
Enable GPU support for Zeta and Polygamma ops.
PiperOrigin-RevId: 211082683
| -rw-r--r-- | tensorflow/core/kernels/cwise_op_zeta.cc | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/cwise_op_zeta.cc b/tensorflow/core/kernels/cwise_op_zeta.cc index 2c5538534c..dc064eec5f 100644 --- a/tensorflow/core/kernels/cwise_op_zeta.cc +++ b/tensorflow/core/kernels/cwise_op_zeta.cc @@ -18,4 +18,9 @@ limitations under the License. namespace tensorflow { REGISTER2(BinaryOp, CPU, "Zeta", functor::zeta, float, double); REGISTER2(BinaryOp, CPU, "Polygamma", functor::polygamma, float, double); + +#if GOOGLE_CUDA +REGISTER2(BinaryOp, GPU, "Zeta", functor::zeta, float, double); +REGISTER2(BinaryOp, GPU, "Polygamma", functor::polygamma, float, double); +#endif } // namespace tensorflow |