examples/research_projects/distillation/train.py/freeze_token_type_embeddings
def freeze_token_type_embeddings(student, args):
    if args.student_type == "roberta":
        student.roberta.embeddings.token_type_embeddings.weight.requires_grad = False