fix json export to remove middleware

neuralmagic · Feb 20, 2024 · ccdd9ea · ccdd9ea
1 parent 1eed6f6
commit ccdd9ea
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/deepsparse/benchmark/benchmark_pipeline.py b/src/deepsparse/benchmark/benchmark_pipeline.py
@@ -346,6 +346,11 @@ def _clear_measurements():
             "Generated no batch timings, try extending benchmark time with '--time'"
         )
 
+    if SupportedTasks.is_text_generation(task) or SupportedTasks.is_code_generation(
+        task
+    ):
+        kwargs.pop("middleware_manager")
+
     return batch_times, total_run_time, num_streams