Update snippet slightly

README.md

@@ -60,20 +60,20 @@ Using `AutoModelForMaskedLM`:
 ```python
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
-checkpoint = "answerdotai/ModernBERT-base"
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForMaskedLM.from_pretrained(checkpoint)
+model_id = "answerdotai/ModernBERT-base"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForMaskedLM.from_pretrained(model_id)
 
 text = "The capital of France is [MASK]."
 inputs = tokenizer(text, return_tensors="pt")
 outputs = model(**inputs)
 
 # To get predictions for the mask:
-logits = outputs.logits
-masked_index = (inputs["input_ids"] == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
-predicted_token_id = logits[0, masked_index].argmax(axis=-1)
+masked_index = inputs["input_ids"][0].tolist().index(tokenizer.mask_token_id)
+predicted_token_id = outputs.logits[0, masked_index].argmax(axis=-1)
 predicted_token = tokenizer.decode(predicted_token_id)
 print("Predicted token:", predicted_token)
+# Predicted token:  Paris
 ```
 
 Using a pipeline: