[TVMC] Switch profile flag to use new profiler (apache#8710)

jiangjiajun · Aug 13, 2021 · ccc09fa · ccc09fa
1 parent 395b308
commit ccc09fa
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git 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
@@ -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)