Skip to content
This repository was archived by the owner on Feb 3, 2025. It is now read-only.

Commit ccc606f

Browse files
author
DEKHTIARJonathan
committed
[Benchmarking-Py] Adding converter.save() before exiting with TF_TRT_BENCHMARK_EARLY_QUIT=1
1 parent 5622106 commit ccc606f

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

tftrt/benchmarking-python/benchmark_runner.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,6 +367,12 @@ def engine_build_input_fn(num_batches, model_phase):
367367
pass
368368

369369
if strtobool(os.environ.get("TF_TRT_BENCHMARK_EARLY_QUIT", "0")):
370+
# Save the result if needed
371+
converter.save(self._args.output_saved_model_dir)
372+
logging.info(
373+
f"Converted graph saved to "
374+
f"`{self._args.output_saved_model_dir}`"
375+
)
370376
sys.exit(0)
371377

372378
if self._args.optimize_offline:

0 commit comments

Comments
 (0)