Using hyperparameter-search in Trainer

Is it at all possible to pass in the model itself to the hp_space? 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