from torch.cuda.amp import autocast
from transformers import AutoModelWithLMHead
from torch.nn import Module


class GPT2LMHeadModel(Module):
    def __init__(self, use_amp: bool = False):
        super(GPT2LMHeadModel, self).__init__()
        self.use_amp = use_amp
        self.transformer = AutoModelWithLMHead.from_pretrained("distilgpt2")

    def forward(self, *args, **kwargs):
        with autocast(self.use_amp):
            return self.transformer.forward(*args, **kwargs)