diff options
| -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 |