| #include <torch/library.h> |
|
|
| #include "registration.h" |
| #include "torch_binding.h" |
|
|
| TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| ops.def("viterbi(Tensor! path, Tensor! score, Tensor! bp, Tensor emissions, Tensor transitions, Tensor priors) -> ()"); |
| ops.def("dtw(Tensor! path, Tensor! path_len, Tensor! D, Tensor cost) -> ()"); |
| ops.def("ctc_forced_align(Tensor! frames, Tensor! score, Tensor log_probs, Tensor targets, int blank) -> ()"); |
|
|
| #if defined(CUDA_KERNEL) || defined(ROCM_KERNEL) |
| ops.impl("viterbi", torch::kCUDA, &viterbi); |
| ops.impl("dtw", torch::kCUDA, &dtw); |
| ops.impl("ctc_forced_align", torch::kCUDA, &ctc_forced_align); |
| #endif |
| } |
|
|
| REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |
|
|
