Update benchmark_utils.py

liuzh47 · Aug 10, 2020 · a43e65b · a43e65b
1 parent 13db82f
commit a43e65b
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/scripts/benchmarks/benchmark_utils.py b/scripts/benchmarks/benchmark_utils.py
@@ -699,12 +699,8 @@ def run_forward():
                 out.wait_to_read()
 
         timeit.repeat(run_forward, repeat=1, number=3)
-        # mx_profiler.set_state('run')
-        # run_forward()
-        # mx_profiler.set_state('stop')
-        # print(mx_profiler.dumps(reset=True))
         runtimes = timeit.repeat(run_forward, repeat=self._repeat, number=3)
-
+        mxnet.npx.waitall()
         # Profile memory
         if self._use_gpu:
             nvml.nvmlInit()