fill-mask usage from transformers pipeline - nlp

Question

I fine-tune a gpt2 language model and I am generation the text according to my model by using following lines of code:
generator = pipeline('text-generation', tokenizer='gpt2', model='data/out')
print(generator('Once upon a time', max_length=40)[0]['generated_text'])
Now I want to do the prediction of only next word with the probabilities. I know we can do it by using 'fill-mask' but I don't know how to do it. When I put 'fill-mask' inplace of 'text-generation', I am getting this error:
"Unrecognized configuration class <class 'transformers.models.gpt2.configuration_gpt2.GPT2Config'> for this kind of AutoModel: AutoModelForMaskedLM.
Model type should be one of BigBirdConfig, Wav2Vec2Config, ConvBertConfig, LayoutLMConfig, DistilBertConfig, AlbertConfig, BartConfig, MBartConfig, CamembertConfig, XLMRobertaConfig, LongformerConfig, RobertaConfig, SqueezeBertConfig, BertConfig, MobileBertConfig, FlaubertConfig, XLMConfig, ElectraConfig, ReformerConfig, FunnelConfig, MPNetConfig, TapasConfig, DebertaConfig, DebertaV2Config, IBertConfig.".
generator = pipeline('fill-mask', tokenizer='gpt2', model='data/out') // this line is giving me the above mentioned error.
Please let me know how can I fix this issue. Any kind of help would be greatly appreciated.
Thanks in advance.
The whole code for better understanding.
from transformers import (
GPT2Tokenizer,
DataCollatorForLanguageModeling,
TextDataset,
GPT2LMHeadModel,
TrainingArguments,
Trainer,
pipeline)
train_path = 'parsed_data.txt'
test_path = 'parsed_data.txt'
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
train_dataset = TextDataset(tokenizer=tokenizer, file_path=train_path, block_size=128)
test_dataset = TextDataset(tokenizer=tokenizer, file_path=test_path, block_size=128)
model = GPT2LMHeadModel.from_pretrained('gpt2')
training_args = TrainingArguments(output_dir = 'data/out', overwrite_output_dir = True, per_device_train_batch_size = 32, per_device_eval_batch_size = 32, learning_rate = 5e-5, num_train_epochs = 3,)
trainer = Trainer(model = model, args = training_args, data_collator=data_collator, train_dataset = train_dataset, eval_dataset = test_dataset)
trainer.train()
trainer.save_model()
generator = pipeline('fill-mask', tokenizer='gpt2', model='data/out')