diff --git a/main.py b/main.py
index 5c0f991..041c0cc 100644
--- a/main.py
+++ b/main.py
@@ -68,7 +68,7 @@ def predict(context_tokens):
 def generate(prompt):
     context_tokens = enc.encode(prompt)
     pred = predict(context_tokens)
-    pred = pred[0][len(context_tokens):]
+    pred = pred[0]["output"][len(context_tokens):]
     output = enc.decode(pred)
     return output