diff --git a/single_shot_inference.py b/single_shot_inference.py
index ef384ef8..09d48998 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -503,6 +503,7 @@ def main():
     
     # ==== Decode: generate new tokens ====
     print(f"\nDecoding (max {MAX_NEW_TOKENS} new tokens)...")
+    all_tokens = generated.copy()  # track full sequence including prompt
     
     for step in range(MAX_NEW_TOKENS):
         t0 = time.time()