padding_cuda.h 450 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
6
7
8
9
10
11
12
#pragma once

#include <torch/extension.h>

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
           torch::Tensor>
padded_index_cuda(torch::Tensor rowptr, torch::Tensor rowcount,
                  torch::Tensor binptr);

torch::Tensor padded_index_select_cuda(torch::Tensor src, torch::Tensor col,
                                       torch::Tensor index,
                                       torch::Tensor fill_value);