Update README.md

README.md

@@ -61,13 +61,13 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 question = f'{tokenizer.cls_token}Where do I live?'  # '<cls>Where do I live?'
 context = 'My name is Sarah and I live in London'
 encoding = tokenizer(question, context, return_tensors="pt")
-start_scores, end_scores = model(
+start_scores, end_scores, _ = model(
   encoding["input_ids"],
   attention_mask=encoding["attention_mask"],
   return_dict=False
 )
 
-all_tokens = tokenizer.convert_ids_to_tokens(input_ids[0].tolist())
+all_tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"][0].tolist())
 answer_tokens = all_tokens[torch.argmax(start_scores):torch.argmax(end_scores) + 1]
 answer = tokenizer.decode(tokenizer.convert_tokens_to_ids(answer_tokens))
 # 'London'