unexpected keyword argument 'mask'

[anuragdalia]

keep getting this issue when trying to train with
python train.py --config-name full_configs/gpt2_next_token_mlm

File "/home/user/SuperTinyLanguageModels/trainers/base_trainer.py", line 234, in run_training_loop      
   losses, perplexities = self.estimate_performance()                                                         
 File "/opt/conda/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context     
   return func(*args, **kwargs)                                                                               
 File "/home/user/SuperTinyLanguageModels/trainers/base_trainer.py", line 133, in estimate_performance   
   losses[i] = self.loss_fn(output, y, mask=mask)                                                             
TypeError: next_token_mlm_loss_fn() got an unexpected keyword argument 'mask'