Update README.md

README.md

@@ -20,9 +20,9 @@ We hope that the community especially researchers and professionals engaged in t
 Here is how to use this model to get the features of a given text in PyTorch:
 
 ```python
-from transformers import BertTokenizer, BertModel, AutoTokenizer, AutoModel
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') # AutoTokenizer.from_pretrained('bert-base-uncased')
-model = BertModel.from_pretrained("vasugoel/K-12BERT") # AutoModel.from_pretrained('bert-base-uncased')
+from transformers import BertTokenizer, BertModel, AutoTokenizer, AutoModelForMaskedLM
+tokenizer = BertTokenizer.from_pretrained('vasugoel/K-12BERT') # AutoTokenizer.from_pretrained('vasugoel/K-12BERT')
+model = BertModel.from_pretrained("vasugoel/K-12BERT") # AutoModelForMaskedLM.from_pretrained('vasugoel/K-12BERT')
 text = "Replace me by any text you'd like."
 encoded_input = tokenizer(text, return_tensors='pt')
 output = model(**encoded_input)