pytorch/torch/csrc/jit/script/init.cpp
2018-02-07 20:43:45 +01:00

21 lines
523 B
C++

#include "torch/csrc/jit/script/init.h"
#include "torch/csrc/jit/script/compiler.h"
namespace torch {
namespace jit {
namespace script {
void initJitScriptBindings(PyObject* module) {
auto m = py::handle(module).cast<py::module>();
py::class_<CompilationUnit>(m, "CompilationUnit")
.def(
"get_graph",
&CompilationUnit::getGraph,
py::return_value_policy::reference);
m.def("_jit_script_compile", jitScriptCompile);
}
} // namespace script
} // namespace jit
} // namespace torch