提交 eb66681f 编写于 作者: M Megvii Engine Team 提交者: huangxinda

fix(mge/random): fix delete_rng_handel

GitOrigin-RevId: a836e13c2385dbfb5df1a31529a457cbfaaa36f5
上级 d6db4fea
......@@ -467,7 +467,9 @@ void init_ops(py::module m) {
m.def("delete_rng_handle", [](size_t handle){
// RNG op might execute after handle released due to async dispatch, so
// we need sync before delete a handle to avoid memory leak or use-after-free
python::interpreter_for_py->sync();
if(python::interpreter_for_py->check_available()){
python::interpreter_for_py->sync();
}
mgb::CompNode::sync_all();
py_task_q.wait_all_task_finish();
rng::delete_handle(handle);
......
......@@ -60,6 +60,7 @@ struct ChannelImpl : Interpreter::Channel {
DeviceTensorND get_dev_tensor(Handle) override;
bool check_available() override;
void sync() override;
void close() override;
......@@ -112,8 +113,6 @@ private:
const SmallVector<LogicalTensorDesc>& input_descs,
SmallVector<Handle>* outputs);
bool check_available();
void push_scope(std::string, State&);
void pop_scope(std::string, State&);
......
......@@ -41,6 +41,7 @@ struct Interpreter {
virtual DeviceTensorND get_dev_tensor(Handle) = 0;
virtual bool check_available() = 0;
virtual void sync() = 0;
virtual void close() = 0;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册