from transformers import GPT2Model, GPT2LMHeadModel from .configuration_fake_custom import FakeCustomConfig class FakeCustomModel(GPT2Model): _auto_class = "AutoModel" config_class = FakeCustomConfig class FakeCustomModelForCausalLM(GPT2LMHeadModel): _auto_class = "AutoModelForCausalLM" config_class = FakeCustomConfig