mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
* Rename autograd namespace to torch and change torch.h into python.h * Include torch.h instead of python.h in test/cpp/api * Change some mentions of torch.h to python.h in C++ extensions * Set paths directly, without find_path
21 lines
457 B
C++
21 lines
457 B
C++
#include <torch/python.h>
|
|
|
|
#include "doubler.h"
|
|
|
|
using namespace at;
|
|
|
|
Tensor exp_add(Tensor x, Tensor y);
|
|
|
|
Tensor tanh_add(Tensor x, Tensor y) {
|
|
return x.tanh() + y.tanh();
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("tanh_add", &tanh_add, "tanh(x) + tanh(y)");
|
|
m.def("exp_add", &exp_add, "exp(x) + exp(y)");
|
|
py::class_<Doubler>(m, "Doubler")
|
|
.def(py::init<int, int>())
|
|
.def("forward", &Doubler::forward)
|
|
.def("get", &Doubler::get);
|
|
}
|