Update README.md

README.md

@@ -68,7 +68,7 @@ example = tokenizer(example, return_tensors='pt', padding=False).to(device) # to
 with torch.no_grad():
     logits = model(**example).logits.cpu().detach() # get logits from model
 
-probability = F.softmax(output, dim=-1) # use softmax to get "confidence" in the prediction
+probability = F.softmax(logits, dim=-1) # use softmax to get "confidence" in the prediction
 prediction = probability.argmax(dim=-1) # 0 for no interaction, 1 for interaction
 ```