Update README.md

README.md

@@ -48,10 +48,17 @@ inputs = tokenizer(
 with torch.no_grad():
     outputs = model(**inputs)
     logits = outputs.logits
-    pred_class = torch.argmax(logits, dim=-1)
+    probs = torch.sigmoid(logits)
+    pred_mask = probs > 0.5
 
-print("logits shape:", logits.shape)      # [batch_size, num_labels]
-print("pred classes:", pred_class.tolist())
+print("logits:", logits)  
+print("logits shape:", logits.shape)  
+print("probs over 0.5:", probs > 0.5)      # [batch_size, num_labels]
+print("pred label mask:", pred_mask.tolist())
+print(
+    "pred label indices:",
+    [[i for i, on in enumerate(row) if on] for row in pred_mask.tolist()],
+)
 ```