diff --git a/python/tvm/driver/tvmc/runner.py b/python/tvm/driver/tvmc/runner.py index 99bce0d3c0e0e..7fbb7ef87344a 100644 --- a/python/tvm/driver/tvmc/runner.py +++ b/python/tvm/driver/tvmc/runner.py @@ -17,6 +17,7 @@ """ Provides support to run compiled networks both locally and remotely. """ +from contextlib import ExitStack import json import logging import pathlib @@ -487,7 +488,10 @@ def run_module( # Remote RPC (running on a micro target) logger.debug("Running on remote RPC (micro target).") try: - session = tvm.micro.Session(project_.transport()).open() + with ExitStack() as stack: + session = tvm.micro.Session(project_.transport()) + # Enter context manager 'session' so it's active from now on. + stack.enter_context(session) except: raise TVMCException("Could not open a session with the micro target.") else: