.

app.py

@@ -39,15 +39,18 @@ def predict_product(num1, num2):
 
     generated_ids = inputs['input_ids']
     past_key_values = None
-    for _ in range(MAX_PRODUCT_DIGITS):  # Set a maximum limit to prevent infinite loops
-        outputs = model.generate(
-            input_ids=generated_ids,
-            max_new_tokens=1,
-            do_sample=False,
-            past_key_values=past_key_values,
-            return_dict_in_generate=True,
-            use_cache=True
-        )
+    for step in range(MAX_PRODUCT_DIGITS):  # Set a maximum limit to prevent infinite loops
+        generation_kwargs = {
+            'input_ids': generated_ids,
+            'max_new_tokens': 1,
+            'do_sample': False,
+            'past_key_values': past_key_values,
+            'return_dict_in_generate': True,
+            'use_cache': True
+        }
+        if step == 0:
+            del generation_kwargs['past_key_values']
+        outputs = model.generate(**generation_kwargs)
         generated_ids = outputs.sequences
         next_token_id = generated_ids[0, -1]
         print (next_token_id)