tweak

app.py

@@ -20,8 +20,8 @@ result = model(batch)
 # first indice is neutral, second is toxic
 prediction = nn.functional.softmax(result.logits, dim=-1)
 
-neutralProb = round(prediction[0][0], 4)
-toxicProb = round(prediction[0][1], 4)
+neutralProb = round(prediction.data[0], 4)
+toxicProb = round(prediction.data[1], 4)
 
 print("Classification Probabilities")
 print(f"Neutral: {neutralProb}")