Actual labels for Optuna hyperparameter importance plot

Question

I am using Optuna as HPO in my Pytorch program to train a NN. After training, I generate the Contour and the parameter importance visualization, however when Optuna generates this p[lot only identifies each parameter as 0,1,2,3,4 but I need to know the actual label of each, for example 0=Learning Rate, 1=Batch Size, 2=Number of epochs, ... etc.

This is the segment of code for the objective function, and the one that handles the visualization:

def optimize_hyperparameters(trial, X, y, device, args):
    # Define the hyperparameters to be tuned
    batch_size = trial.suggest_categorical('batch_size', [ 32, 64, 128, 256, 512, 1024 ])
    lr = trial.suggest_float('lr', 1e-5, 1e-1, log=True)
    epochs = trial.suggest_int('epochs', 10, 100)
    folds = trial.suggest_int('folds', 3, 10)  # You can adjust the range as needed
    processors = trial.suggest_int('processors', 2, 8)  # Adjust this list based on your requirements
    # Update args with the suggested values
    args.batch_size = batch_size
    args.lr = lr
    args.epochs = epochs
    args.folds = folds  # Updating the folds
    args.processors = processors  # Updating the processors

    # Use the updated args in the training function
    results, _ = parallel_kfold(X, y, device, args)
    # Measure the start time
    start_time = time.time()
##
    processing_time = time.time() - start_time

    accuracies = [result[1] for result in results]
    average_accuracy = sum(accuracies) / len(accuracies)

    if average_accuracy < 0.90:
        raise optuna.exceptions.TrialPruned()

    # Composite score: e.g., accuracy - (penalty_factor * processing_time)
    penalty_factor = 0.001  # Adjust based on how much you want to penalize time
    composite_score = average_accuracy - (penalty_factor * processing_time)

    return composite_score

def main():
...

    study = optuna.create_study(
    direction='maximize',
    study_name='myresearch_study',   
    storage='postgresql://XXXXXXX:dbpass@localhost/optuna_study',load_if_exists=True)

    n_trials = args.trials ## set num of trials
    study.optimize(lambda trial: optimize_hyperparameters(trial, X_train, y_train, device, args), n_trials)  

    print("Best hyperparameters: ", study.best_params)

    plt.switch_backend('agg')

    # Plotting optimization history
    ax = plot_optimization_history(study)
    fig = ax.figure
    fig.set_size_inches(8, 6)   
    fig.savefig("optimization_history.png")

    # Plotting parameter importances
    ax = plot_param_importances(study)
    fig = ax.figure
    fig.set_size_inches(8, 6)  # Set the dimensions after getting the figure
    fig.savefig("param_importances.png")
...

I tried the previous code to generate the hyperparameter importance plot

I expected a plot with the actual hyperparameters label, not 0,1,2,3,4 which I don't know what they are.