Pre-fit ModelListGP and pass it to Models.BOTORCH_MODULAR in Ax Client about ax HOT 3 OPEN

def create_gp_model_kwargs():
    return {'surrogate' : Surrogate(botorch_model_class = train_default_gp_model()),

            'botorch_acqf_class': qNoisyExpectedHypervolumeImprovement}

def train_default_gp_model(self):

    train_x, train_y = initialize_training_data()

    train_x = normalize(X=train_x_motors, bounds=params_bounds_tensor)

    models = []

    for i in range(train_y.shape[-1]):

        models.append(SingleTaskGP(train_X = train_x,

                                   train_Y = train_y[:, i].unsqueeze(-1),

                                   outcome_transform = Standardize(m = 1)))

        model = ModelListGP(*models)

        mll = SumMarginalLogLikelihood(model.likelihood, model)

        fit_gpytorch_mll(mll)

    return model

parameters_list = [dict(name = param,
type = 'range',
bounds = range_values,
value_type = 'float') 
                   for param in all_parameters]
objectives = {obj: ObjectiveProperties(minimize=True, threshold=hypervolume_reference_point) 
              for obj in all_objectives}
steps = [GenerationStep(model = Models.SOBOL, # Add 1 sobol to avoid "ax.exceptions.core.DataRequiredError: StandardizeY transform requires non-empty data."

                        num_trials = 1),

         GenerationStep(model = Models.BOTORCH_MODULAR,

                        num_trials. = num_bo_trials,

                        model_kwargs = create_gp_model_kwargs())]

# Create Ax Client and optimize:

ax_client = AxClient(generation_strategy = GenerationStrategy(steps=steps),

                     random_seed = fixed_seed,

                     verbose_logging = True)

ax_client.create_experiment(name = 'Ax_Modular',

                            parameters = parameters_list,

                            objectives = objectives,

                            overwrite_existing_experiment = True)
for i in range(BO_trials):

    suggested_values, trial_index = ax_client.get_next_trial()
    # Throw the error after one Sobol trial
    loss_values = evaluate_suggested_values(suggested_values=suggested_values)
 
    ax_client.complete_trial(trial_index=trial_index, raw_data=loss_values.copy())