Update README.md

README.md

@@ -30,7 +30,7 @@ with torch.no_grad():
 
 # Get predictions
 threshold = 0.5
-probs = torch.sigmoid(outputs.logits)
+probs = torch.sigmoid(outputs.logits).squeeze()
 predictions = [model.config.id2label[i] for i, prob in enumerate(probs) if prob > threshold]
 print(predictions)
 # -> ['ppr_pos', 'wtp_neg'] (patient-provider/staff relationships:positive, waiting time for appointment:negative)