Problem with pest synthetic problem

[TimShinners]

I have been trying to use the pest synthetic problem in the MCBO package. If I try to run

from HEBO.MCBO.mcbo import task_factory
from HEBO.MCBO.mcbo.optimizers.bo_builder import BO_ALGOS

task = task_factory(task_name="pest")
search_space = task.get_search_space()

optimizer = BO_ALGOS['CoCaBO'].build_bo(search_space=search_space, n_init=10)

for i in range(100):
    x = optimizer.suggest(1)
    y = task(x)
    optimizer.observe(x, y)
    print(f'Iteration {i + 1:3d}/{100:3d} - f(x) = {y[0, 0]:.3f} - f(x*) = {optimizer.best_y:.3f}')

then it leads to an assertion error on line 101 in pest.py: assert x.ndim == 1 and len(x) == self._n_stages

I can get around the error by manually setting the number of stages by saying task._n_stages=24. I noticed on line 108 of pest.py it says for i in range(1, n_stages): which might have something to do with it.

Am I doing something wrong or is there potentially a bug in the pest problem?

[Antoine Grosnit]

Thanks for spotting this., you are right, I have now fixed this 0e30d83