File size: 649 Bytes
6fc683c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
/**
* Copyright 2017-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the license found in the
* LICENSE file in the root directory of this source tree.
*/
#include "alignment_train_cuda.h"
#include "utils.h"
namespace {
void alignmentTrainCUDA(
const torch::Tensor& p_choose,
torch::Tensor& alpha,
float eps) {
CHECK_INPUT(p_choose);
CHECK_INPUT(alpha);
alignmentTrainCUDAWrapper(p_choose, alpha, eps);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def(
"alignment_train_cuda",
&alignmentTrainCUDA,
"expected_alignment_from_p_choose (CUDA)");
}
} // namespace
|