diff --git a/python/tvm/driver/tvmc/model.py b/python/tvm/driver/tvmc/model.py index 7dc3fd4cdd36..a9516e1e2c42 100644 --- a/python/tvm/driver/tvmc/model.py +++ b/python/tvm/driver/tvmc/model.py @@ -371,7 +371,7 @@ def import_package(self, package_path: str): class TVMCResult(object): """A class that stores the results of tvmc.run and provides helper utilities.""" - def __init__(self, outputs: Dict[str, np.ndarray], times: List[str]): + def __init__(self, outputs: Dict[str, np.ndarray], times: List[float]): """Create a convenience wrapper around the output of tvmc.run Parameters diff --git a/python/tvm/driver/tvmc/runner.py b/python/tvm/driver/tvmc/runner.py index 916139874579..8515bc9b053c 100644 --- a/python/tvm/driver/tvmc/runner.py +++ b/python/tvm/driver/tvmc/runner.py @@ -417,7 +417,9 @@ def run_module( # Run must be called explicitly if profiling if profile: logger.info("Running the module with profiling enabled.") - module.run() + report = module.profile() + # This print is intentional + print(report) # create the module time evaluator (returns a function) timer = module.module.time_evaluator("run", dev, number=number, repeat=repeat)