Update main.py

main.py

@@ -97,7 +97,7 @@ def predict(text: str) -> str:
         pred = torch.argmax(out, dim=-1)[0]
 
     # Keep the prediction same length as input
-    pred = pred[:input_len]
+    pred = pred[:input_len+1]
 
     return "".join(idx_to_char[i.item()] for i in pred)