diff --git a/ax/modelbridge/strategies/alebo.py b/ax/modelbridge/strategies/alebo.py index d6032e03575..4cb324a8c96 100644 --- a/ax/modelbridge/strategies/alebo.py +++ b/ax/modelbridge/strategies/alebo.py @@ -16,18 +16,23 @@ from ax.modelbridge.random import RandomModelBridge from ax.modelbridge.torch import TorchModelBridge from ax.modelbridge.transforms.centered_unit_x import CenteredUnitX +from ax.modelbridge.transforms.int_to_float import IntToFloat +from ax.modelbridge.transforms.remove_fixed import RemoveFixed from ax.modelbridge.transforms.standardize_y import StandardizeY from ax.models.random.alebo_initializer import ALEBOInitializer from ax.models.torch.alebo import ALEBO +ALEBO_X_trans = [RemoveFixed, IntToFloat, CenteredUnitX] + + def get_ALEBOInitializer( search_space: SearchSpace, B: np.ndarray, **model_kwargs: Any ) -> RandomModelBridge: return RandomModelBridge( search_space=search_space, model=ALEBOInitializer(B=B, **model_kwargs), - transforms=[CenteredUnitX], + transforms=ALEBO_X_trans, # pyre-ignore ) @@ -45,7 +50,7 @@ def get_ALEBO( search_space=search_space, data=data, model=ALEBO(B=B, **model_kwargs), - transforms=[CenteredUnitX, StandardizeY], + transforms=ALEBO_X_trans + [StandardizeY], # pyre-ignore torch_dtype=B.dtype, torch_device=B.device, )