Update app.py

app.py

@@ -23,10 +23,10 @@ classes = {
 def prediction(tweet, model, tokenizer):
     inputs = tokenizer(tweet, return_tensors="pt", padding=True, truncation=True)
     outputs = model(**inputs)
-    predicted_class = torch.argmax(outputs.logits, dim=1)
+    predicted_class = torch.argmax(outputs.logits, dim=1).item()  # convert to scalar integer
     predicted_prob = torch.softmax(outputs.logits, dim=1)[0][predicted_class].item()
     return classes[predicted_class], predicted_prob
-    
+
 def create_table(predictions):
     data = {'Tweet': [], 'Highest Toxicity Class': [], 'Probability': []}
     for tweet, prediction in predictions.items():