Passing in model to optuna

Hello all! One of the things I would like to see is which pretrained model from huggingface works best for my dataset. I would like to pass in the different model names as one of the parameters to optimize on to Optuna. How can I integrate that with the HuggingFace Trainer?

For example if you want to test out different models such as “base-base-uncased” or “roberta-base” or “mpnet-base-v2” etc.

I currently have it set up like this:

epochs = trial.suggest_categorical("epochs", EPOCHS)
    batch_size = trial.suggest_categorical("batch_size", BATCH_SIZE)
    learning_rate = trial.suggest_categorical("learning_rate", LEARNING_RATES)
    scheduler = trial.suggest_categorical("scheduler", SCHEDULERS)
    model_name = trial.suggest_categorical("model_name", MODEL_NAMES)
    hp_space = {
        "model_name": model_name,
        "batch_size": batch_size,
        "learning_rate": learning_rate,
        "scheduler": scheduler,
        "epochs": epochs,

I am not sure how to pass this in correctly to Trainer

Modifying the model_init function as below should help

MODEL_NAMES=[ “base-base-uncased” , “roberta-base” , “mpnet-base-v2”]
def model_init(params):
    if params is not None:
    return your_model(model_name)