def load_gpt2(model_name="gpt2"):
    """
    load original gpt2 and save off for quicker loading

    Args:
        model_name: GPT-2

    Returns:
        GPT-2 model

    """

    model = GPT2LMHeadModel.from_pretrained(model_name, output_hidden_states=True)
    torch.save(model.state_dict(), model_name + "local.pt")
    return model

    load original gpt2 and save off for quicker loading

    Args:
        model_name: GPT-2

    Returns:
        GPT-2 model

    
dpkd/transformers/examples/research_projects/information-gain-filtration/igf/igf.py