Update app.py

app.py

@@ -15,8 +15,8 @@ story = open('story.txt').read().replace("\n\n", "\n").replace("\n", " ").strip(
 
 def ask(question):
     inputs = tokenizer(story, question, truncation=True, padding=True)
-    prediction = torch.softmax(model(**inputs), dim=-1).squeeze()
-    return [{c: prediction[i].item() for c, i in CLASSES}]
+    prediction = torch.softmax(model(**inputs).logits, dim=-1).squeeze()
+    return {c: prediction[i].item() for c, i in CLASSES}
 
 
 demo = gr.Interface(