mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
a0118533ef
Additionally: - add support for calling functions that are not methods in the Python frontend - add an end-to-end test for the Python frontend - add a capture_stdout helper for checking that `print` actually works
21 lines
584 B
C++
21 lines
584 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(py::init<>())
|
|
.def("get_graph", &CompilationUnit::getGraph,
|
|
py::return_value_policy::reference)
|
|
.def("define_function", &CompilationUnit::defineFunction);
|
|
m.def("_jit_script_compile", jitScriptCompile);
|
|
}
|
|
|
|
} // namespace script
|
|
} // namespace jit
|
|
} // namespace torch
|