forked from jiuyuan/InfiniTensor
feat: 导出分配内存和运行推理的接口
Signed-off-by: YdrMaster <ydrml@hotmail.com>
This commit is contained in:
parent
d7e52054e6
commit
6871fff02b
|
@ -76,6 +76,12 @@ class GraphHandlerObj {
|
|||
const optional<vector<int>> &steps);
|
||||
Tensor pad(Tensor input, Tensor output, const vector<int> &pads,
|
||||
const optional<vector<int>> &axes);
|
||||
|
||||
//------ runtime
|
||||
|
||||
inline void data_malloc() { g->dataMalloc(); }
|
||||
|
||||
inline void run() { g->getRuntime()->run(g); }
|
||||
};
|
||||
|
||||
} // namespace infini
|
||||
|
|
|
@ -102,7 +102,9 @@ void init_graph_builder(py::module &m) {
|
|||
.def("pad",
|
||||
py::overload_cast<Tensor, Tensor, const vector<int> &,
|
||||
const optional<vector<int>> &>(&Handler::pad),
|
||||
policy::move);
|
||||
policy::move)
|
||||
.def("data_malloc", &Handler::data_malloc, policy::automatic)
|
||||
.def("run", &Handler::run, policy::automatic);
|
||||
}
|
||||
|
||||
} // namespace infini
|
||||
|
|
Loading…
Reference in New Issue