p35189024
/
InfiniTensor
forked from jiuyuan/InfiniTensor
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

feat: 导出数据 

Signed-off-by: YdrMaster <ydrml@hotmail.com>
This commit is contained in:
YdrMaster 2023-09-18 16:57:41 +08:00
parent 04c2725312
commit 8e4e392a49
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/ffi/ffi_infinitensor.cc
View File

@ -52,6 +52,9 @@ class Handler {
graph->getRuntime()->run(graph);
#endif
}
template <class T> std::vector<T> copyout(size_t i) {
return _outputs[i]->copyout<T>();
}
};
using TExport = std::tuple<Name, int, std::vector<std::variant<Name, int>>>;
@ -220,6 +223,7 @@ void register_refactor(py::module &m) {
.def("fill_edge_info", &Handler::fillEdgeInfo)
.def("substitute", &Handler::substitute)
.def("run_cuda", &Handler::runCuda)
.def("copy_out_float", &Handler::copyout<float>)
.def("set_input", &Handler::setInput);
py::class_<NodeExport>(m, "NodeExport")
.def(py::init<std::shared_ptr<Handler>>())