diff --git a/src/models/input_ids.cpp b/src/models/input_ids.cpp
index 4cba33fe1..93de5f8ab 100644
--- a/src/models/input_ids.cpp
+++ b/src/models/input_ids.cpp
@@ -76,6 +76,7 @@ void DefaultInputIDs::Update(DeviceSpan<int32_t>& new_tokens) {
 #endif
     }
 
+    is_prompt_ = false;
     state_.inputs_[input_index_] = value_.get();
     return;
   }