Update README.md

README.md

@@ -25,8 +25,7 @@ model = RobertaForSequenceClassification.from_pretrained('s-nlp/roberta_first_to
 batch = tokenizer.encode("You are amazing!", return_tensors="pt")
 
 output = model(batch)
-predicted_label = torch.sigmoid(output.logits).argmax().item()
-# 0 for neutral, 1 for toxic
+# idx 0 for neutral, idx 1 for toxic
 ```
 
 ## Citation