10 #include "torch/script.h"
14 const torch::Tensor& row_splits);
16 #ifdef BUILD_CUDA_MODULE
18 torch::Tensor ReduceSubarraysSumCUDA(
const torch::Tensor& values,
19 const torch::Tensor& row_splits);
torch::Tensor ReduceSubarraysSumCPU(const torch::Tensor &values, const torch::Tensor &row_splits)
Definition: ReduceSubarraysSumOpKernel.cpp:16