Make change in benchmark as well

microsoft · May 4, 2024 · 4ee082f · 4ee082f
1 parent 0b7cec5
commit 4ee082f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/benchmark/python/benchmark_e2e.py b/benchmark/python/benchmark_e2e.py
@@ -85,6 +85,9 @@ def main(args):
             generator.generate_next_token()
         if args.print_model_output: print(tokenizer.decode(generator.get_sequence(0)))
 
+        # Delete the generator to free the captured graph for the next generator, if graph capture is enabled
+        del generator
+
     tokenize_times = []
     prompt_times = []
     token_gen_times = []
@@ -141,6 +144,9 @@ def main(args):
         wall_clock_times.append(wall_clock_end_time - wall_clock_start_time)
         if args.print_model_output: print(tokenizer.decode(generator.get_sequence(0)))
 
+        # Delete the generator to free the captured graph for the next generator, if graph capture is enabled
+        del generator
+
     # Calculate tokenization metrics
     avg_tokenization_latency_s = sum(tokenize_times) / len(tokenize_times)
     avg_tokenization_latency_ms = avg_tokenization_latency_s * 1000