.

app.py

@@ -40,15 +40,15 @@ 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(
+        outputs = model.generate(
             input_ids=generated_ids,
+            max_new_tokens=1,
             past_key_values=past_key_values,
+            return_dict_in_generate=True,
             use_cache=True
         )
-        logits = outputs.logits
-        
-        next_token_id = torch.argmax(logits[:, -1, :], dim=-1)
-        generated_ids = torch.cat((generated_ids, next_token_id.view(1,-1)), dim=-1)
+        generated_ids = outputs.sequences
+        next_token_id = generated_ids[0, -1]
         print (next_token_id)
         
         if next_token_id.item() == eos_token_id: