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);
|
const optional<vector<int>> &steps);
|
||||||
Tensor pad(Tensor input, Tensor output, const vector<int> &pads,
|
Tensor pad(Tensor input, Tensor output, const vector<int> &pads,
|
||||||
const optional<vector<int>> &axes);
|
const optional<vector<int>> &axes);
|
||||||
|
|
||||||
|
//------ runtime
|
||||||
|
|
||||||
|
inline void data_malloc() { g->dataMalloc(); }
|
||||||
|
|
||||||
|
inline void run() { g->getRuntime()->run(g); }
|
||||||
};
|
};
|
||||||
|
|
||||||
} // namespace infini
|
} // namespace infini
|
||||||
|
|
|
@ -102,7 +102,9 @@ void init_graph_builder(py::module &m) {
|
||||||
.def("pad",
|
.def("pad",
|
||||||
py::overload_cast<Tensor, Tensor, const vector<int> &,
|
py::overload_cast<Tensor, Tensor, const vector<int> &,
|
||||||
const optional<vector<int>> &>(&Handler::pad),
|
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
|
} // namespace infini
|
||||||
|
|
Loading…
Reference in New Issue