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