Update README.md

README.md

@@ -29,12 +29,10 @@ import torch
 sentence = "Let's have a [MASK]."
 
 model.eval()
-encoded_inputs = tokenizer([sentence], padding='longest')
-input_ids = torch.tensor(encoded_inputs['input_ids'])
-attention_mask = torch.tensor(encoded_inputs['attention_mask'])
-output = model(input_ids, attention_mask=attention_mask)
+inputs = tokenizer([sentence], padding='longest', return_tensors='pt')
+output = model(inputs['input_ids'], attention_mask=inputs['attention_mask'])
 
-mask_index = input_ids.tolist()[0].index(103)
+mask_index = inputs['input_ids'].tolist()[0].index(103)
 masked_token = output['logits'][0][mask_index].argmax(axis=-1)
 predicted_token = tokenizer.decode(masked_token)