Update README.md

README.md

@@ -31,7 +31,7 @@ hypothesis = "The movie was good."
 input = tokenizer(premise, hypothesis, truncation=True, return_tensors="pt")
 output = model(input["input_ids"].to(device))  # device = "cuda:0" or "cpu"
 prediction = torch.softmax(output["logits"][0], -1).tolist()
-label_names = ["entailment", "neutral", "contradiction"]
+label_names = ["entailment", "not_entailment"]
 prediction = {name: round(float(pred) * 100, 1) for pred, name in zip(prediction, label_names)}
 print(prediction)
 ```