Update app.py

app.py

@@ -41,14 +41,13 @@ model = AutoModelForSequenceClassification.from_pretrained(modelPath)
 # RoBERTA: [0]: neutral, [1]: toxic
 encoding = tokenizer.encode(txt, return_tensors='pt')
 result = model(encoding)
-
-result
-if (result.logits.size(dim=1) >= 2):
-    pad = (0, 1)
-    result = nn.functional.pad(result, pad, "constant", 0)
 result
 
 # transform logit to get probabilities
+if (result.logits.size(dim=1) < 2):
+    pad = (0, 1)
+    result = nn.functional.pad(result, pad, "constant", 0)
+    prediction = nn.functional.softmax(result, dim=-1)
 prediction = nn.functional.softmax(result.logits, dim=-1)
 prediction
 #neutralProb = prediction.data[0][neutralIndex]