tweak

app.py

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