                   framework::Tensor* output,
                   framework::Tensor* summed_ids = nullptr,
                   const std::string pooltype = "SUM") {
+    if (pooltype == "MEAN") {
+      // Sum the segment id num first
+      T DimTileSize = 8;
+      auto input_length_size = segment_ids.numel();
+      auto total_stripe_count =
+          (input_length_size + DimTileSize - 1) / DimTileSize;
+      auto config = platform::GetGpuLaunchConfig1D(ctx, total_stripe_count);
+      SegmentSumIdsKernel<
+          T, IndexT, IndexT(8)><<<config.block_per_grid.x,
+                                  config.thread_per_block.x, 0, ctx.stream()>>>(
+          segment_ids.data<IndexT>(), summed_ids->data<T>(), input_length_size,
+          total_stripe_count);
+    }
+
     auto h = ArrangeHelper<IndexT>(input.numel(), segment_ids.dims()[0],
                                    output->dims()[0]);
     auto config = platform::GetGpuLaunchConfig1D(ctx, h.total_stripe_count);
     if (pooltype == "MEAN") {
-      SegmentMeanCustomKernel<
+      SegmentMeanKernel<
           T, IndexT, IndexT(8)><<<config.block_per_grid.x,
                                   config.thread_per_block.x, 0, ctx.stream()>>>(
           segment_ids.data<IndexT>(), input.data<T>(), output->data<T>(),

using Tensor = framework::Tensor;

template <typename T, typename Index, int DimTileSize>
__global__ void SegmentMeanCustomKernel(
    const Index* segment_ids, const T* input, T* output, T* summed_ids,
    const Index input_length_size, const Index inner_dim_size,
    const Index output_length_size, const Index total_stripe_count) {
  CUDA_KERNEL_LOOP(stripe_index, total_stripe_count) {
    const Index segment_offset = stripe_index % inner_dim_size;
    const Index dim_index_base =

    if (dim_index_base > 0) {
      last_segment_id = segment_ids[dim_index_base - 1];
    }
    T sum = T(0);
    for (Index j = 0; j < actual_height; j++) {
      Index current_segment_id = segment_ids[dim_index_base + j];
      if (current_segment_id > last_segment_id) {
        // reset the interval value which do not have corresponding ids.
        for (Index interval_id = last_segment_id + 1;
             interval_id < current_segment_id; ++interval_id) {
          *(output + interval_id * inner_dim_size + segment_offset) = T(0);
        }

        if (j > 0) {
          Index output_index =
              last_segment_id * inner_dim_size + segment_offset;

          if (last_segment_id == first_segment_id) {
            platform::CudaAtomicAdd(output + output_index,
                                    sum / *(summed_ids + last_segment_id));
          } else {
            *(output + output_index) = sum / *(summed_ids + last_segment_id);
          }
          sum = T(0);
        }
      }
      sum += input[(dim_index_base + j) * inner_dim_size + segment_offset];
      last_segment_id = current_segment_id;
    }
    Index output_index = last_segment_id * inner_dim_size + segment_offset;
    platform::CudaAtomicAdd(output + output_index,
                            sum / *(summed_ids + last_segment_id));
  }
