Update app.py

app.py

@@ -66,7 +66,7 @@ if st.button('Submit'):
         encoding = tokenizer(text, return_tensors="pt")
         outputs = model(**encoding)
         predictions = outputs.logits.argmax(-1)
-        number = predictions.cpu().detach().numpy()
+        number = predictions[0].cpu().detach().numpy()
         probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
         fig = plt.figure()