diff --git a/src/runtime/rpc/rpc_session.cc b/src/runtime/rpc/rpc_session.cc index 2c38294b967d91e2e60195a25b7ad902ca4fb040..923194cf4f1b87c4e25015db881a42a30aa640b2 100644 --- a/src/runtime/rpc/rpc_session.cc +++ b/src/runtime/rpc/rpc_session.cc @@ -783,6 +783,9 @@ void RPCSession::Shutdown() { void RPCSession::ServerLoop() { std::lock_guard<std::recursive_mutex> lock(mutex_); + if (const auto* f = Registry::Get("tvm.contrib.rpc.server.start")) { + (*f)(); + } TVMRetValue rv; CHECK(HandleUntilReturnEvent(&rv, false, nullptr) == RPCCode::kShutdown); LOG(INFO) << "Shutdown...";