File size: 655 Bytes
bfdbd33 0a16898 bfdbd33 0a16898 bfdbd33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | #include <torch/library.h>
#include "registration.h"
#include "torch_binding.h"
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("physarum_step(Tensor! trail, Tensor! tmp, Tensor! ax, Tensor! ay, "
"Tensor! ah, int seed, float sense_dist, float sense_ang, float turn, "
"float speed, float deposit, float decay) -> ()");
ops.def("flow_cg(Tensor cE, Tensor cS, Tensor b, Tensor! p, int ground, "
"int max_iters, float tol) -> ()");
#if defined(CPU_KERNEL)
ops.impl("physarum_step", torch::kCPU, &physarum_step);
ops.impl("flow_cg", torch::kCPU, &flow_cg);
#endif
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|