Hi @prajjwal1, I think you should be able to do this by inherit  from GPT2PreTrainedModel instead of nn.module , and if all you want to do is replace the loss function then you can take the code for GPT2LMHeadModel and just replace the loss calculation part.