From f55c321397895c60db063938a2bc76b4f08ede38 Mon Sep 17 00:00:00 2001
From: prajwalc22 <prajwalchoudhari22@gmail.com>
Date: Thu, 17 Apr 2025 10:15:21 +0530
Subject: [PATCH] Address review feedback: Fix prefill_tbatch_size and variable
 placement issues

---
 gemma/run.cc | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/gemma/run.cc b/gemma/run.cc
index 36d1bc2..56cdb75 100644
--- a/gemma/run.cc
+++ b/gemma/run.cc
@@ -156,8 +156,7 @@ void ReplGemma(const ThreadingArgs& threading, const InferenceArgs& inference,
     std::cout << token_text << std::flush;
     return true;
   };
-  // Flag to check if we should exit after processing non-interactive prompt
-  bool exit_after_generation = !inference.prompt.empty();
+
   while (true) {  // Loop until user quits.
     tokens_generated_this_turn = 0;
 
@@ -224,7 +223,7 @@ void ReplGemma(const ThreadingArgs& threading, const InferenceArgs& inference,
     std::cout << "\n\n";
 
     // Break the loop if in non-interactive mode
-    if (exit_after_generation) {
+    if (!inference.prompt.empty()) {
       break;
     }