diff --git a/imperative/python/src/ops.cpp b/imperative/python/src/ops.cpp index e151fadb..22ed862e 100644 --- a/imperative/python/src/ops.cpp +++ b/imperative/python/src/ops.cpp @@ -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); diff --git a/imperative/src/impl/interpreter/interpreter_impl.h b/imperative/src/impl/interpreter/interpreter_impl.h index 00fcf8e8..14b3b2de 100644 --- a/imperative/src/impl/interpreter/interpreter_impl.h +++ b/imperative/src/impl/interpreter/interpreter_impl.h @@ -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& input_descs, SmallVector* outputs); - bool check_available(); - void push_scope(std::string, State&); void pop_scope(std::string, State&); diff --git a/imperative/src/include/megbrain/imperative/interpreter.h b/imperative/src/include/megbrain/imperative/interpreter.h index a4f29ff4..f7f5028f 100644 --- a/imperative/src/include/megbrain/imperative/interpreter.h +++ b/imperative/src/include/megbrain/imperative/interpreter.h @@ -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;