@@ -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 { | |||||
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: | |||||
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 { | |||||
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; | ||||