From ec2c339b53dd3a5cd76041b1bb0166db4e26465c Mon Sep 17 00:00:00 2001 From: Davide Fiocco Date: Wed, 28 Nov 2018 14:53:46 +0100 Subject: [PATCH] Updated quick-start example with `BertForMaskedLM` As `convert_ids_to_tokens` returns a list, the code in the README currently throws an `AssertionError`, so I propose I quick fix. --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index bf5453853197..eb337d8253f4 100644 --- a/README.md +++ b/README.md @@ -142,7 +142,7 @@ predictions = model(tokens_tensor, segments_tensors) # confirm we were able to predict 'henson' predicted_index = torch.argmax(predictions[0, masked_index]).item() -predicted_token = tokenizer.convert_ids_to_tokens([predicted_index]) +predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])[0] assert predicted_token == 'henson' ```