Yw22's picture
init demo
d711508
raw
history blame
815 Bytes
#include <torch/torch.h>
#include <vector>
#include <iostream>
#include <torch/extension.h>
std::vector<at::Tensor> forward_fast_block_diag_cuda(
at::Tensor input);
std::vector<at::Tensor> forward_fast_block_diag(
at::Tensor input
) {
return forward_fast_block_diag_cuda(input);
}
std::vector<at::Tensor> backward_fast_block_diag_cuda(
at::Tensor grad_output,
at::Tensor input);
std::vector<at::Tensor> backward_fast_block_diag(
at::Tensor grad_output,
at::Tensor input
) {
return backward_fast_block_diag_cuda(grad_output, input);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward_fast_block_diag, "FAST BLOCK DIAG (CUDA)");
m.def("backward", &backward_fast_block_diag, "FAST BLOCK DIAG backward (CUDA)");
}