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 = np.round(torch.softmax(model(**inputs).logits, 1).squeeze().numpy(), 3)
-    return {c: prediction[i] for c, i in CLASSES}
+    prediction = torch.round(torch.softmax(model(**inputs).logits, 1).squeeze().numpy(), 3)
+    return {c: prediction[i].item() for c, i in CLASSES}
 
 
 demo = gr.Interface(