[Benchmarking-Py] Adding converter.save() before exiting with `TF_T…

…RT_BENCHMARK_EARLY_QUIT=1`
tensorflow · Sep 22, 2022 · ccc606f · ccc606f
1 parent 5622106
commit ccc606f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tftrt/benchmarking-python/benchmark_runner.py b/tftrt/benchmarking-python/benchmark_runner.py
@@ -367,6 +367,12 @@ def engine_build_input_fn(num_batches, model_phase):
                 pass
 
             if strtobool(os.environ.get("TF_TRT_BENCHMARK_EARLY_QUIT", "0")):
+                # Save the result if needed
+                converter.save(self._args.output_saved_model_dir)
+                logging.info(
+                    f"Converted graph saved to "
+                    f"`{self._args.output_saved_model_dir}`"
+                )
                 sys.exit(0)
 
             if self._args.optimize_offline: