diff options
Diffstat (limited to 'tensorflow/core/kernels/segment_reduction_ops.cc')
| -rw-r--r-- | tensorflow/core/kernels/segment_reduction_ops.cc | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/tensorflow/core/kernels/segment_reduction_ops.cc b/tensorflow/core/kernels/segment_reduction_ops.cc index c87ce78e05..2328fc6afd 100644 --- a/tensorflow/core/kernels/segment_reduction_ops.cc +++ b/tensorflow/core/kernels/segment_reduction_ops.cc @@ -320,7 +320,9 @@ class SegmentSumGPUOp : public AsyncOpKernel { REGISTER_CPU_KERNEL_SEGMENT("SegmentSum", Eigen::internal::SumReducer<type>, \ type, index_type, 0); \ REGISTER_CPU_KERNEL_SEGMENT( \ - "SegmentProd", Eigen::internal::ProdReducer<type>, type, index_type, 1) + "SegmentMean", Eigen::internal::MeanReducer<type>, type, index_type, 0); \ + REGISTER_CPU_KERNEL_SEGMENT( \ + "SegmentProd", Eigen::internal::ProdReducer<type>, type, index_type, 1); #define REGISTER_REAL_CPU_KERNELS_ALL(type) \ REGISTER_REAL_CPU_KERNELS(type, int32); \ |