Update README.md

README.md

@@ -28,7 +28,7 @@ WWM ensures that all subword tokens of a selected word are masked together, enco
 
 - **Developed by:** [Mohammad Mahdi Heydari Asl / infocube]  
 - **Funded by [optional]:** [More Information Needed]
-- **Shared by:** [[HYDARIM7](https://huggingface.co/InfocubeSrl/LexCube)]  
+- **Shared by:** [[HYDARIM7](https://huggingface.co/HYDARIM7)]  
 - **Model type:** Transformer, BERT-based Masked Language Model  
 - **Language(s) (NLP):** Italian  
 - **License:** Apache-2.0  
@@ -71,9 +71,11 @@ text = "La legge [MASK] approvata dal parlamento."
 inputs = tokenizer(text, return_tensors="pt")
 outputs = model(**inputs)
 
-mask_token_index = torch.where(inputs["input_ids"][0] == tokenizer.mask_token_id)[0]
-predicted_token_id = outputs.logits[0, mask_token_index].argmax(axis=-1)
-print("Prediction:", tokenizer.decode(predicted_token_id))
+mask_index = (inputs["input_ids"][0] == tokenizer.mask_token_id).nonzero()[0]
+predicted_id = outputs.logits[0, mask_index].argmax()
+predicted_token = tokenizer.decode(predicted_id)
+
+print("Prediction:", predicted_token)
 ```