Update README.md

README.md

@@ -61,14 +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(
+output = model(
   encoding["input_ids"],
-  attention_mask=encoding["attention_mask"],
-  return_dict=False
+  attention_mask=encoding["attention_mask"]
 )
 
 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_tokens = all_tokens[torch.argmax(output["start_logits"]):torch.argmax(output["end_logits"]) + 1]
 answer = tokenizer.decode(tokenizer.convert_tokens_to_ids(answer_tokens))
 # 'London'
 ```