diff --git a/mlforecast/_modidx.py b/mlforecast/_modidx.py
index a802a118..d8627a8c 100644
--- a/mlforecast/_modidx.py
+++ b/mlforecast/_modidx.py
@@ -5,15 +5,35 @@
                 'doc_host': 'https://Nixtla.github.io',
                 'git_url': 'https://github.com/Nixtla/mlforecast',
                 'lib_path': 'mlforecast'},
-  'syms': { 'mlforecast.auto': { 'mlforecast.auto.AutoMLForecast': ('auto.html#automlforecast', 'mlforecast/auto.py'),
+  'syms': { 'mlforecast.auto': { 'mlforecast.auto.AutoCatboost': ('auto.html#autocatboost', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoCatboost.__init__': ('auto.html#autocatboost.__init__', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoElasticNet': ('auto.html#autoelasticnet', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoElasticNet.__init__': ('auto.html#autoelasticnet.__init__', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoLasso': ('auto.html#autolasso', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoLasso.__init__': ('auto.html#autolasso.__init__', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoLightGBM': ('auto.html#autolightgbm', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoLightGBM.__init__': ('auto.html#autolightgbm.__init__', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoLinearRegression': ('auto.html#autolinearregression', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoLinearRegression.__init__': ( 'auto.html#autolinearregression.__init__',
+                                                                                    'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoMLForecast': ('auto.html#automlforecast', 'mlforecast/auto.py'),
                                  'mlforecast.auto.AutoMLForecast.__init__': ('auto.html#automlforecast.__init__', 'mlforecast/auto.py'),
                                  'mlforecast.auto.AutoMLForecast._seasonality_based_config': ( 'auto.html#automlforecast._seasonality_based_config',
                                                                                                'mlforecast/auto.py'),
                                  'mlforecast.auto.AutoMLForecast.fit': ('auto.html#automlforecast.fit', 'mlforecast/auto.py'),
                                  'mlforecast.auto.AutoMLForecast.predict': ('auto.html#automlforecast.predict', 'mlforecast/auto.py'),
                                  'mlforecast.auto.AutoMLForecast.save': ('auto.html#automlforecast.save', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoModel': ('auto.html#automodel', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoModel.__init__': ('auto.html#automodel.__init__', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoModel.__repr__': ('auto.html#automodel.__repr__', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoRandomForest': ('auto.html#autorandomforest', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoRandomForest.__init__': ('auto.html#autorandomforest.__init__', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoRidge': ('auto.html#autoridge', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoRidge.__init__': ('auto.html#autoridge.__init__', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoXGBoost': ('auto.html#autoxgboost', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.AutoXGBoost.__init__': ('auto.html#autoxgboost.__init__', 'mlforecast/auto.py'),
                                  'mlforecast.auto.catboost_space': ('auto.html#catboost_space', 'mlforecast/auto.py'),
-                                 'mlforecast.auto.elasticnet_space': ('auto.html#elasticnet_space', 'mlforecast/auto.py'),
+                                 'mlforecast.auto.elastic_net_space': ('auto.html#elastic_net_space', 'mlforecast/auto.py'),
                                  'mlforecast.auto.lasso_space': ('auto.html#lasso_space', 'mlforecast/auto.py'),
                                  'mlforecast.auto.lightgbm_space': ('auto.html#lightgbm_space', 'mlforecast/auto.py'),
                                  'mlforecast.auto.linear_regression_space': ('auto.html#linear_regression_space', 'mlforecast/auto.py'),
diff --git a/mlforecast/auto.py b/mlforecast/auto.py
index 71a845f1..be91d4bd 100644
--- a/mlforecast/auto.py
+++ b/mlforecast/auto.py
@@ -2,31 +2,26 @@
 
 # %% auto 0
 __all__ = ['lightgbm_space', 'xgboost_space', 'catboost_space', 'linear_regression_space', 'ridge_space', 'lasso_space',
-           'elasticnet_space', 'random_forest_space', 'AutoMLForecast']
+           'elastic_net_space', 'random_forest_space', 'AutoModel', 'AutoLightGBM', 'AutoXGBoost', 'AutoCatboost',
+           'AutoLinearRegression', 'AutoRidge', 'AutoLasso', 'AutoElasticNet', 'AutoRandomForest', 'AutoMLForecast']
 
 # %% ../nbs/auto.ipynb 2
-from collections import defaultdict
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import numpy as np
 import optuna
-import pandas as pd
 from sklearn.base import BaseEstimator, clone
-from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import FunctionTransformer
 from utilsforecast.compat import DataFrame
-from utilsforecast.losses import mase, smape
+from utilsforecast.losses import smape
 from utilsforecast.processing import counts_by_id
 from utilsforecast.validation import validate_freq
 
 from . import MLForecast
-from .core import Freq, Models, _get_model_name, _name_models
-from mlforecast.lag_transforms import (
-    ExpandingMean,
-    ExponentiallyWeightedMean,
-    RollingMean,
-)
-from .optimization import mlforecast_objective
+from .core import Freq, _get_model_name, _name_models
+from .lag_transforms import ExponentiallyWeightedMean, RollingMean
+from .optimization import _TrialToConfig, mlforecast_objective
 from mlforecast.target_transforms import (
     Differences,
     LocalStandardScaler,
@@ -34,7 +29,7 @@
 )
 
 # %% ../nbs/auto.ipynb 3
-def lightgbm_space(trial):
+def lightgbm_space(trial: optuna.Trial):
     return {
         "bagging_freq": 1,
         "learning_rate": 0.05,
@@ -49,7 +44,7 @@ def lightgbm_space(trial):
     }
 
 
-def xgboost_space(trial):
+def xgboost_space(trial: optuna.Trial):
     return {
         "n_estimators": trial.suggest_int("n_estimators", 20, 1000),
         "max_depth": trial.suggest_int("max_depth", 1, 10),
@@ -64,8 +59,9 @@ def xgboost_space(trial):
     }
 
 
-def catboost_space(trial):
+def catboost_space(trial: optuna.Trial):
     return {
+        "silent": True,
         "n_estimators": trial.suggest_int("n_estimators", 50, 1000),
         "depth": trial.suggest_int("depth", 1, 10),
         "learning_rate": trial.suggest_float("learning_rate", 1e-3, 0.2, log=True),
@@ -75,25 +71,25 @@ def catboost_space(trial):
     }
 
 
-def linear_regression_space(trial):
+def linear_regression_space(trial: optuna.Trial):
     return {"fit_intercept": trial.suggest_categorical("fit_intercept", [True, False])}
 
 
-def ridge_space(trial):
+def ridge_space(trial: optuna.Trial):
     return {
         "fit_intercept": trial.suggest_categorical("fit_intercept", [True, False]),
         "alpha": trial.suggest_float("alpha", 0.001, 10.0),
     }
 
 
-def lasso_space(trial):
+def lasso_space(trial: optuna.Trial):
     return {
         "fit_intercept": trial.suggest_categorical("fit_intercept", [True, False]),
         "alpha": trial.suggest_float("alpha", 0.001, 10.0),
     }
 
 
-def elasticnet_space(trial):
+def elastic_net_space(trial: optuna.Trial):
     return {
         "fit_intercept": trial.suggest_categorical("fit_intercept", [True, False]),
         "alpha": trial.suggest_float("alpha", 0.001, 10.0),
@@ -101,7 +97,7 @@ def elasticnet_space(trial):
     }
 
 
-def random_forest_space(trial):
+def random_forest_space(trial: optuna.Trial):
     return {
         "n_estimators": trial.suggest_int("n_estimators", 50, 1000),
         "max_depth": trial.suggest_int("max_depth", 1, 10),
@@ -112,56 +108,171 @@ def random_forest_space(trial):
         ),
     }
 
-# %% ../nbs/auto.ipynb 4
-_default_model_spaces = {
-    "LGBMRegressor": lightgbm_space,
-    "XGBRegressor": xgboost_space,
-    "CatBoostRegressor": catboost_space,
-    "LinearRegression": linear_regression_space,
-    "Ridge": ridge_space,
-    "Lasso": lasso_space,
-    "ElasticNet": elasticnet_space,
-    "RandomForest": random_forest_space,
-}
-
-_ModelWithConfig = Tuple[BaseEstimator, Optional[Dict[str, Callable]]]
-
-# %% ../nbs/auto.ipynb 5
+
+class AutoModel:
+    """Structure to hold a model and its search space
+
+    Parameters
+    ----------
+    model : BaseEstimator
+        scikit-learn compatible regressor
+    config : callable
+        function that takes an optuna trial and produces a configuration
+    """
+
+    def __init__(
+        self,
+        model: BaseEstimator,
+        config: _TrialToConfig,
+    ):
+        self.model = model
+        self.config = config
+
+    def __repr__(self):
+        return f"AutoModel(model={_get_model_name(self.model)})"
+
+
+class AutoLightGBM(AutoModel):
+    def __init__(
+        self,
+        config: Optional[_TrialToConfig] = None,
+    ):
+        from mlforecast.compat import LGBMRegressor
+
+        super().__init__(
+            LGBMRegressor(),
+            config if config is not None else lightgbm_space,
+        )
+
+
+class AutoXGBoost(AutoModel):
+    def __init__(
+        self,
+        config: Optional[_TrialToConfig] = None,
+    ):
+        from mlforecast.compat import XGBRegressor
+
+        super().__init__(
+            XGBRegressor(),
+            config if config is not None else xgboost_space,
+        )
+
+
+class AutoCatboost(AutoModel):
+    def __init__(
+        self,
+        config: Optional[_TrialToConfig] = None,
+    ):
+        from mlforecast.compat import CatBoostRegressor
+
+        super().__init__(
+            CatBoostRegressor(),
+            config if config is not None else catboost_space,
+        )
+
+
+class AutoLinearRegression(AutoModel):
+    def __init__(
+        self,
+        config: Optional[_TrialToConfig] = None,
+    ):
+        from sklearn.linear_model import LinearRegression
+
+        super().__init__(
+            LinearRegression(),
+            config if config is not None else linear_regression_space,
+        )
+
+
+class AutoRidge(AutoModel):
+    def __init__(
+        self,
+        config: Optional[_TrialToConfig] = None,
+    ):
+        from sklearn.linear_model import Ridge
+
+        super().__init__(
+            Ridge(),
+            config if config is not None else ridge_space,
+        )
+
+
+class AutoLasso(AutoModel):
+    def __init__(
+        self,
+        config: Optional[_TrialToConfig] = None,
+    ):
+        from sklearn.linear_model import Lasso
+
+        super().__init__(
+            Lasso(),
+            config if config is not None else lasso_space,
+        )
+
+
+class AutoElasticNet(AutoModel):
+    def __init__(
+        self,
+        config: Optional[_TrialToConfig] = None,
+    ):
+        from sklearn.linear_model import ElasticNet
+
+        super().__init__(
+            ElasticNet(),
+            config if config is not None else elastic_net_space,
+        )
+
+
+class AutoRandomForest(AutoModel):
+    def __init__(
+        self,
+        config: Optional[_TrialToConfig] = None,
+    ):
+        from sklearn.ensemble import RandomForestRegressor
+
+        super().__init__(
+            RandomForestRegressor(),
+            config if config is not None else random_forest_space,
+        )
+
+# %% ../nbs/auto.ipynb 6
 class AutoMLForecast:
+    """Hyperparameter optimization helper
+
+    Parameters
+    ----------
+    models : list or dict
+        Auto models to be optimized.
+    freq : str or int
+        pandas' or polars' offset alias or integer denoting the frequency of the series.
+    season_length : int
+        Length of the seasonal period. This is used for producing the feature space.
+    init_config : callable, optional (default=None)
+        Function that takes an optuna trial and produces a configuration passed to the MLForecast constructor.
+    fit_config : callable, optional (default=None)
+        Function that takes an optuna trial and produces a configuration passed to the MLForecast fit method.
+    num_threads : int (default=1)
+        Number of threads to use when computing the features.
+    """
 
     def __init__(
         self,
-        models_with_configs: Union[List[_ModelWithConfig], Dict[str, _ModelWithConfig]],
+        models: Union[List[AutoModel], Dict[str, AutoModel]],
         freq: Freq,
         season_length: int,
-        init_config: Optional[Callable] = None,
-        fit_config: Optional[Callable] = None,
+        init_config: Optional[_TrialToConfig] = None,
+        fit_config: Optional[_TrialToConfig] = None,
         num_threads: int = 1,
     ):
         self.freq = freq
         self.season_length = season_length
         self.num_threads = num_threads
-        if isinstance(models_with_configs, list):
-            names = _name_models([_get_model_name(m) for m, _ in models_with_configs])
-            self.models_with_configs = {
-                name: (model, config)
-                for name, (model, config) in zip(names, models_with_configs)
-            }
-        elif isinstance(models_with_configs, dict):
-            self.models_with_configs = models_with_configs
+        if isinstance(models, list):
+            model_names = _name_models([_get_model_name(m) for m in models])
+            models_with_names = dict(zip(model_names, models))
         else:
-            raise ValueError(
-                "`models_with_configs` should be a list of tuples "
-                "or a dict from str to tuple."
-            )
-        for name, (model, config) in self.models_with_configs.items():
-            if config is not None:
-                continue
-            if name not in _default_model_spaces:
-                raise NotImplementedError(
-                    f"{name} does not have a default config. Please provide one."
-                )
-            self.models_with_configs[name] = (model, _default_model_spaces[name])
+            models_with_names = models
+        self.models = models_with_names
         self.init_config = init_config
         if fit_config is not None:
             self.fit_config = fit_config
@@ -173,7 +284,7 @@ def _seasonality_based_config(
         h: int,
         min_samples: int,
         min_value: float,
-    ) -> Callable:
+    ) -> _TrialToConfig:
         # target transforms
         candidate_targ_tfms = [
             None,
@@ -205,9 +316,7 @@ def _seasonality_based_config(
                 )
 
         # lags
-        candidate_lags = [None]
-        if self.season_length > 1:
-            candidate_lags.append([self.season_length])
+        candidate_lags = [None, [self.season_length]]
         seasonality2extra_candidate_lags = {
             7: [
                 [7, 14],
@@ -218,6 +327,9 @@ def _seasonality_based_config(
                 range(1, 25),
                 range(24, 24 * 7 + 1, 24),
             ],
+            52: [
+                range(4, 53, 4),
+            ],
         }
         candidate_lags.extend(
             seasonality2extra_candidate_lags.get(self.season_length, [])
@@ -265,13 +377,9 @@ def _seasonality_based_config(
             52: ["week", "year"],
             60: ["weekday", "hour", "second"],
         }
-        date_features = seasonality2date_features.get(self.season_length, None)
+        candidate_date_features = seasonality2date_features.get(self.season_length, [])
         if isinstance(self.freq, int):
-            date_features = None
-        if date_features is not None:
-            use_date_features = trial.suggest_int("use_date_features", 0, 1)
-            if not use_date_features:
-                date_features = None
+            candidate_date_features = []
 
         def config(trial):
             # target transforms
@@ -293,6 +401,16 @@ def config(trial):
             else:
                 lag_transforms = None
 
+            # date features
+            if candidate_date_features:
+                use_date_features = trial.suggest_int("use_date_features", 0, 1)
+                if use_date_features:
+                    date_features = candidate_date_features
+                else:
+                    date_features = None
+            else:
+                date_features = None
+
             return {
                 "lags": lags,
                 "target_transforms": target_transforms,
@@ -308,13 +426,45 @@ def fit(
         n_windows: int,
         h: int,
         num_samples: int,
-        loss: Optional[Callable] = None,
+        loss: Optional[Callable[[DataFrame, DataFrame], float]] = None,
         id_col: str = "unique_id",
         time_col: str = "ds",
         target_col: str = "y",
         study_kwargs: Optional[Dict[str, Any]] = None,
         optimize_kwargs: Optional[Dict[str, Any]] = None,
     ) -> "AutoMLForecast":
+        """Carry out the optimization process.
+        Each model is optimized independently and the best one is trained on all data
+
+        Parameters
+        ----------
+        df : pandas or polars DataFrame
+            Series data in long format.
+        n_windows : int
+            Number of windows to evaluate.
+        h : int
+            Forecast horizon.
+        num_samples : int
+            Number of trials to run
+        loss : callable, optional (default=None)
+            Function that takes the validation and train dataframes and produces a float.
+            If `None` will use the average SMAPE across series.
+        id_col : str (default='unique_id')
+            Column that identifies each serie.
+        time_col : str (default='ds')
+            Column that identifies each timestep, its values can be timestamps or integers.
+        target_col : str (default='y')
+            Column that contains the target.
+        study_kwargs : dict, optional (default=None)
+            Keyword arguments to be passed to the optuna.Study constructor.
+        optimize_kwargs : dict, optional (default=None)
+            Keyword arguments to be passed to the Study.optimize method.
+
+        Returns
+        -------
+        AutoMLForecast
+            object with best models and optimization results
+        """
         validate_freq(df[time_col], self.freq)
         if self.init_config is not None:
             init_config = self.init_config
@@ -337,18 +487,14 @@ def loss(df, train_df):
         if "sampler" not in study_kwargs:
             # for reproducibility
             study_kwargs["sampler"] = optuna.samplers.TPESampler(seed=0)
-        if optimize_kwargs is None:
-            optimize_kwargs = {}
-        if "n_jobs" not in optimize_kwargs:
-            optimize_kwargs["n_jobs"] = 1
 
         self.results_ = []
         self.models_ = {}
-        for name, (model, model_config) in self.models_with_configs.items():
+        for name, auto_model in self.models.items():
 
             def config_fn(trial: optuna.Trial) -> float:
                 return {
-                    "model_params": model_config(trial),
+                    "model_params": auto_model.config(trial),
                     "mlf_init_params": {
                         **init_config(trial),
                         "num_threads": self.num_threads,
@@ -359,8 +505,8 @@ def config_fn(trial: optuna.Trial) -> float:
             objective = mlforecast_objective(
                 df=df,
                 config_fn=config_fn,
-                eval_fn=loss,
-                model=model,
+                loss=loss,
+                model=auto_model.model,
                 freq=self.freq,
                 n_windows=n_windows,
                 h=h,
@@ -372,7 +518,7 @@ def config_fn(trial: optuna.Trial) -> float:
             study.optimize(objective, n_trials=num_samples, **optimize_kwargs)
             self.results_.append(study)
             best_config = study.best_trial.user_attrs["config"]
-            best_model = clone(model)
+            best_model = clone(auto_model.model)
             best_model.set_params(**best_config["model_params"])
             self.models_[name] = MLForecast(
                 models={name: best_model},
@@ -387,6 +533,20 @@ def predict(
         h: int,
         X_df: Optional[DataFrame] = None,
     ) -> DataFrame:
+        """ "Compute forecasts
+
+        Parameters
+        ----------
+        h : int
+            Number of periods to predict.
+        X_df : pandas or polars DataFrame, optional (default=None)
+            Dataframe with the future exogenous features. Should have the id column and the time column.
+
+        Returns
+        -------
+        pandas or polars DataFrame
+            Predictions for each serie and timestep, with one column per model.
+        """
         all_preds = None
         for name, model in self.models_.items():
             preds = model.predict(h=h, X_df=X_df)
@@ -396,6 +556,12 @@ def predict(
                 all_preds[name] = preds[name]
         return all_preds
 
-    def save(self, path: str) -> None:
+    def save(self, path: Union[str, Path]) -> None:
+        """Save MLForecast objects
+
+        Parameters
+        ----------
+        path : str or pathlib.Path
+            Directory where artifacts will be stored."""
         for name, model in self.models_.items():
             model.save(f"{path}/{name}")
diff --git a/mlforecast/compat.py b/mlforecast/compat.py
index 2cfe2bfd..dd29519a 100644
--- a/mlforecast/compat.py
+++ b/mlforecast/compat.py
@@ -6,10 +6,34 @@
 # %% ../nbs/compat.ipynb 1
 try:
     from catboost import CatBoostRegressor
+except ImportError:
+
+    class CatBoostRegressor:
+        def __init__(self, *args, **kwargs):
+            raise ImportError("Please install catboost to use this model.")
+
+
+try:
+    from lightgbm import LGBMRegressor
+except ImportError:
+
+    class LGBMRegressor:
+        def __init__(self, *args, **kwargs):
+            raise ImportError("Please install lightgbm to use this model.")
+
+
+try:
     from window_ops.shift import shift_array
 except ImportError:
 
     def shift_array(*_args, **_kwargs):
         raise Exception
 
-    class CatBoostRegressor: ...
+
+try:
+    from xgboost import XGBRegressor
+except ImportError:
+
+    class XGBRegressor:
+        def __init__(self, *args, **kwargs):
+            raise ImportError("Please install xgboost to use this model.")
diff --git a/mlforecast/optimization.py b/mlforecast/optimization.py
index 4cbf50b9..3e4006bb 100644
--- a/mlforecast/optimization.py
+++ b/mlforecast/optimization.py
@@ -5,24 +5,26 @@
 
 # %% ../nbs/optimization.ipynb 2
 import copy
-from typing import Callable, List, Optional
+from typing import Any, Callable, Dict, Optional
 
 import numpy as np
 import optuna
 import utilsforecast.processing as ufp
 from sklearn.base import BaseEstimator, clone
 from utilsforecast.compat import DataFrame
-from utilsforecast.losses import smape
 
 from . import MLForecast
 from .compat import CatBoostRegressor
 from .core import Freq
 
 # %% ../nbs/optimization.ipynb 3
+_TrialToConfig = Callable[[optuna.Trial], Dict[str, Any]]
+
+# %% ../nbs/optimization.ipynb 4
 def mlforecast_objective(
     df: DataFrame,
-    config_fn: Callable,
-    eval_fn: Callable,
+    config_fn: _TrialToConfig,
+    loss: Callable[[DataFrame, DataFrame], float],
     model: BaseEstimator,
     freq: Freq,
     n_windows: int,
@@ -30,7 +32,35 @@ def mlforecast_objective(
     id_col: str = "unique_id",
     time_col: str = "ds",
     target_col: str = "y",
-) -> Callable:
+) -> _TrialToConfig:
+    """optuna objective function for the MLForecast class
+
+    Parameters
+    ----------
+    config_fn : callable
+        Function that takes an optuna trial and produces a configuration with the following keys:
+        - model_params
+        - mlf_init_params
+        - mlf_fit_params
+    loss : callable
+        Function that takes the validation and train dataframes and produces a float.
+    model : BaseEstimator
+        scikit-learn compatible model to be trained
+    freq : str or int
+        pandas' or polars' offset alias or integer denoting the frequency of the series.
+    n_windows : int
+        Number of windows to evaluate.
+    h : int
+        Forecast horizon.
+    id_col : str (default='unique_id')
+        Column that identifies each serie.
+    time_col : str (default='ds')
+        Column that identifies each timestep, its values can be timestamps or integers.
+    target_col : str (default='y')
+        Column that contains the target.
+    study_kwargs : dict, optional (default=None)
+    """
+
     def objective(trial: optuna.Trial) -> float:
         config = config_fn(trial)
         trial.set_user_attr("config", copy.deepcopy(config))
@@ -95,7 +125,7 @@ def objective(trial: optuna.Trial) -> float:
                     "Please verify that the passed frequency (freq) matches your series' "
                     "and that there aren't any missing periods."
                 )
-            metric = eval_fn(result, train_df=train)
+            metric = loss(result, train_df=train)
             metrics.append(metric)
             trial.report(metric, step=i)
             if trial.should_prune():
diff --git a/nbs/auto.ipynb b/nbs/auto.ipynb
index 8d7080e0..cc851a36 100644
--- a/nbs/auto.ipynb
+++ b/nbs/auto.ipynb
@@ -30,24 +30,22 @@
    "outputs": [],
    "source": [
     "#| export\n",
-    "from collections import defaultdict\n",
-    "from typing import Any, Callable, Dict, List, Optional, Tuple, Union\n",
+    "from pathlib import Path\n",
+    "from typing import Any, Callable, Dict, List, Optional, Union\n",
     "\n",
     "import numpy as np\n",
     "import optuna\n",
-    "import pandas as pd\n",
     "from sklearn.base import BaseEstimator, clone\n",
-    "from sklearn.pipeline import Pipeline\n",
     "from sklearn.preprocessing import FunctionTransformer\n",
     "from utilsforecast.compat import DataFrame\n",
-    "from utilsforecast.losses import mase, smape\n",
+    "from utilsforecast.losses import smape\n",
     "from utilsforecast.processing import counts_by_id\n",
     "from utilsforecast.validation import validate_freq\n",
     "\n",
     "from mlforecast import MLForecast\n",
-    "from mlforecast.core import Freq, Models, _get_model_name, _name_models\n",
-    "from mlforecast.lag_transforms import ExpandingMean, ExponentiallyWeightedMean, RollingMean\n",
-    "from mlforecast.optimization import mlforecast_objective\n",
+    "from mlforecast.core import Freq, _get_model_name, _name_models\n",
+    "from mlforecast.lag_transforms import ExponentiallyWeightedMean, RollingMean\n",
+    "from mlforecast.optimization import _TrialToConfig, mlforecast_objective\n",
     "from mlforecast.target_transforms import Differences, LocalStandardScaler, GlobalSklearnTransformer"
    ]
   },
@@ -59,37 +57,38 @@
    "outputs": [],
    "source": [
     "#| export\n",
-    "def lightgbm_space(trial):\n",
+    "def lightgbm_space(trial: optuna.Trial):\n",
     "    return {\n",
-    "        \"bagging_freq\": 1,\n",
-    "        \"learning_rate\": 0.05,\n",
-    "        \"verbosity\": -1,        \n",
-    "        \"n_estimators\": trial.suggest_int(\"n_estimators\", 20, 1000, log=True),\n",
+    "        'bagging_freq': 1,\n",
+    "        'learning_rate': 0.05,\n",
+    "        'verbosity': -1,        \n",
+    "        'n_estimators': trial.suggest_int('n_estimators', 20, 1000, log=True),\n",
     "        'lambda_l1': trial.suggest_float('lambda_l1', 1e-8, 10.0, log=True),\n",
     "        'lambda_l2': trial.suggest_float('lambda_l2', 1e-8, 10.0, log=True),\n",
     "        'num_leaves': trial.suggest_int('num_leaves', 2, 4096, log=True),\n",
     "        'feature_fraction': trial.suggest_float('feature_fraction', 0.5, 1.0),\n",
     "        'bagging_fraction': trial.suggest_float('bagging_fraction', 0.5, 1.0),\n",
-    "        \"objective\": trial.suggest_categorical(\"objective\", ['l1', 'l2']),\n",
+    "        'objective': trial.suggest_categorical('objective', ['l1', 'l2']),\n",
     "    }\n",
     "\n",
-    "def xgboost_space(trial):\n",
+    "def xgboost_space(trial: optuna.Trial):\n",
     "    return {\n",
-    "        \"n_estimators\": trial.suggest_int(\"n_estimators\", 20, 1000),\n",
+    "        'n_estimators': trial.suggest_int('n_estimators', 20, 1000),\n",
     "        'max_depth': trial.suggest_int('max_depth', 1, 10),\n",
     "        'learning_rate': trial.suggest_float('learning_rate', 1e-3, 0.2, log=True),\n",
     "        'subsample': trial.suggest_float('subsample', 0.1, 1.0),\n",
     "        'bagging_freq': trial.suggest_float('bagging_freq', 0.1, 1.0),\n",
     "        'colsample_bytree': trial.suggest_float('colsample_bytree', 0.1, 1.0),\n",
     "        'min_data_in_leaf': trial.suggest_float('min_data_in_leaf', 1, 100),\n",
-    "        \"reg_lambda\": trial.suggest_float(\"reg_lambda\", 1e-8, 1.0, log=True),\n",
-    "        \"reg_alpha\": trial.suggest_float(\"reg_alpha\", 1e-8, 1.0, log=True),\n",
+    "        'reg_lambda': trial.suggest_float('reg_lambda', 1e-8, 1.0, log=True),\n",
+    "        'reg_alpha': trial.suggest_float('reg_alpha', 1e-8, 1.0, log=True),\n",
     "        'min_child_weight': trial.suggest_int('min_child_weight', 2, 10),\n",
     "    }\n",
     "    \n",
-    "def catboost_space(trial):\n",
+    "def catboost_space(trial: optuna.Trial):\n",
     "    return {\n",
-    "        \"n_estimators\": trial.suggest_int(\"n_estimators\", 50, 1000),\n",
+    "        'silent': True,\n",
+    "        'n_estimators': trial.suggest_int('n_estimators', 50, 1000),\n",
     "        'depth': trial.suggest_int('depth', 1, 10),\n",
     "        'learning_rate': trial.suggest_float('learning_rate', 1e-3, 0.2, log=True),\n",
     "        'subsample': trial.suggest_float('subsample', 0.1, 1.0),\n",
@@ -97,60 +96,210 @@
     "        'min_data_in_leaf': trial.suggest_float('min_data_in_leaf', 1, 100),\n",
     "    }\n",
     "    \n",
-    "def linear_regression_space(trial):\n",
+    "def linear_regression_space(trial: optuna.Trial):\n",
     "    return {\n",
-    "        \"fit_intercept\": trial.suggest_categorical(\"fit_intercept\", [True, False])\n",
+    "        'fit_intercept': trial.suggest_categorical('fit_intercept', [True, False])\n",
     "    }\n",
     "    \n",
-    "def ridge_space(trial):\n",
+    "def ridge_space(trial: optuna.Trial):\n",
     "    return {\n",
-    "        \"fit_intercept\": trial.suggest_categorical(\"fit_intercept\", [True, False]),\n",
+    "        'fit_intercept': trial.suggest_categorical('fit_intercept', [True, False]),\n",
     "        'alpha': trial.suggest_float('alpha', 0.001, 10.0)\n",
     "    }\n",
     "    \n",
-    "def lasso_space(trial):\n",
+    "def lasso_space(trial: optuna.Trial):\n",
     "    return {\n",
-    "        \"fit_intercept\": trial.suggest_categorical(\"fit_intercept\", [True, False]),\n",
+    "        'fit_intercept': trial.suggest_categorical('fit_intercept', [True, False]),\n",
     "        'alpha': trial.suggest_float('alpha', 0.001, 10.0)\n",
     "    }\n",
     "    \n",
-    "def elasticnet_space(trial):\n",
+    "def elastic_net_space(trial: optuna.Trial):\n",
     "    return {\n",
-    "        \"fit_intercept\": trial.suggest_categorical(\"fit_intercept\", [True, False]),\n",
+    "        'fit_intercept': trial.suggest_categorical('fit_intercept', [True, False]),\n",
     "        'alpha': trial.suggest_float('alpha', 0.001, 10.0),\n",
     "        'l1_ratio': trial.suggest_float('l1_ratio', 0.0, 1.0)\n",
     "    }\n",
     "\n",
-    "def random_forest_space(trial):\n",
+    "def random_forest_space(trial: optuna.Trial):\n",
     "    return {\n",
-    "        \"n_estimators\": trial.suggest_int(\"n_estimators\", 50, 1000),\n",
+    "        'n_estimators': trial.suggest_int('n_estimators', 50, 1000),\n",
     "        'max_depth': trial.suggest_int('max_depth', 1, 10),\n",
     "        'min_samples_split': trial.suggest_int('min_child_samples', 1, 100),\n",
     "        'max_features': trial.suggest_float('max_features', 0.5, 1.0),\n",
-    "        \"criterion\": trial.suggest_categorical(\"criterion\", ['squared_error', 'poisson']),\n",
-    "    }"
+    "        'criterion': trial.suggest_categorical('criterion', ['squared_error', 'poisson']),\n",
+    "    }\n",
+    "\n",
+    "class AutoModel:\n",
+    "    \"\"\"Structure to hold a model and its search space\n",
+    " \n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    model : BaseEstimator\n",
+    "        scikit-learn compatible regressor\n",
+    "    config : callable \n",
+    "        function that takes an optuna trial and produces a configuration\n",
+    "    \"\"\"\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        model: BaseEstimator,\n",
+    "        config: _TrialToConfig,\n",
+    "    ):\n",
+    "        self.model = model\n",
+    "        self.config = config\n",
+    "\n",
+    "    def __repr__(self):\n",
+    "        return f'AutoModel(model={_get_model_name(self.model)})'\n",
+    "\n",
+    "class AutoLightGBM(AutoModel):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: Optional[_TrialToConfig] = None,\n",
+    "    ):\n",
+    "        from mlforecast.compat import LGBMRegressor\n",
+    "        super().__init__(\n",
+    "            LGBMRegressor(),\n",
+    "            config if config is not None else lightgbm_space,\n",
+    "        )\n",
+    "\n",
+    "class AutoXGBoost(AutoModel):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: Optional[_TrialToConfig] = None,\n",
+    "    ):\n",
+    "        from mlforecast.compat import XGBRegressor\n",
+    "        super().__init__(\n",
+    "            XGBRegressor(),\n",
+    "            config if config is not None else xgboost_space,\n",
+    "        )\n",
+    "\n",
+    "class AutoCatboost(AutoModel):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: Optional[_TrialToConfig] = None,\n",
+    "    ):\n",
+    "        from mlforecast.compat import CatBoostRegressor\n",
+    "        super().__init__(\n",
+    "            CatBoostRegressor(),\n",
+    "            config if config is not None else catboost_space,\n",
+    "        )\n",
+    "\n",
+    "class AutoLinearRegression(AutoModel):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: Optional[_TrialToConfig] = None,\n",
+    "    ):\n",
+    "        from sklearn.linear_model import LinearRegression\n",
+    "        super().__init__(\n",
+    "            LinearRegression(),\n",
+    "            config if config is not None else linear_regression_space,\n",
+    "        )\n",
+    "\n",
+    "class AutoRidge(AutoModel):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: Optional[_TrialToConfig] = None,\n",
+    "    ):\n",
+    "        from sklearn.linear_model import Ridge\n",
+    "        super().__init__(\n",
+    "            Ridge(),\n",
+    "            config if config is not None else ridge_space,\n",
+    "        )\n",
+    "\n",
+    "class AutoLasso(AutoModel):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: Optional[_TrialToConfig] = None,\n",
+    "    ):\n",
+    "        from sklearn.linear_model import Lasso\n",
+    "        super().__init__(\n",
+    "            Lasso(),\n",
+    "            config if config is not None else lasso_space,\n",
+    "        )\n",
+    "\n",
+    "class AutoElasticNet(AutoModel):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: Optional[_TrialToConfig] = None,\n",
+    "    ):\n",
+    "        from sklearn.linear_model import ElasticNet\n",
+    "        super().__init__(\n",
+    "            ElasticNet(),\n",
+    "            config if config is not None else elastic_net_space,\n",
+    "        )\n",
+    "\n",
+    "class AutoRandomForest(AutoModel):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: Optional[_TrialToConfig] = None,\n",
+    "    ):\n",
+    "        from sklearn.ensemble import RandomForestRegressor\n",
+    "        super().__init__(\n",
+    "            RandomForestRegressor(),\n",
+    "            config if config is not None else random_forest_space,\n",
+    "        )"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c01bfb2f-43ab-4848-add3-3e4c01cf3c49",
+   "id": "b2819ca2",
    "metadata": {},
    "outputs": [],
    "source": [
-    "#| exporti\n",
-    "_default_model_spaces = {\n",
-    "    'LGBMRegressor': lightgbm_space,\n",
-    "    'XGBRegressor': xgboost_space,\n",
-    "    'CatBoostRegressor': catboost_space,\n",
-    "    'LinearRegression': linear_regression_space,\n",
-    "    'Ridge': ridge_space,\n",
-    "    'Lasso': lasso_space,\n",
-    "    'ElasticNet': elasticnet_space,\n",
-    "    'RandomForest': random_forest_space,\n",
-    "}\n",
-    "\n",
-    "_ModelWithConfig = Tuple[BaseEstimator, Optional[Dict[str, Callable]]]"
+    "#| hide\n",
+    "from nbdev import show_doc"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ea9f3693",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "---\n",
+       "\n",
+       "[source](https://github.com/Nixtla/mlforecast/blob/main/mlforecast/auto.py#L115){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
+       "\n",
+       "### AutoModel\n",
+       "\n",
+       ">      AutoModel (model:sklearn.base.BaseEstimator,\n",
+       ">                 config:Callable[[optuna.trial._trial.Trial],Dict[str,Any]])\n",
+       "\n",
+       "Structure to hold a model and its search space\n",
+       "\n",
+       "|    | **Type** | **Details** |\n",
+       "| -- | -------- | ----------- |\n",
+       "| model | BaseEstimator | scikit-learn compatible regressor |\n",
+       "| config | Callable | function that takes an optuna trial and produces a configuration |"
+      ],
+      "text/plain": [
+       "---\n",
+       "\n",
+       "[source](https://github.com/Nixtla/mlforecast/blob/main/mlforecast/auto.py#L115){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
+       "\n",
+       "### AutoModel\n",
+       "\n",
+       ">      AutoModel (model:sklearn.base.BaseEstimator,\n",
+       ">                 config:Callable[[optuna.trial._trial.Trial],Dict[str,Any]])\n",
+       "\n",
+       "Structure to hold a model and its search space\n",
+       "\n",
+       "|    | **Type** | **Details** |\n",
+       "| -- | -------- | ----------- |\n",
+       "| model | BaseEstimator | scikit-learn compatible regressor |\n",
+       "| config | Callable | function that takes an optuna trial and produces a configuration |"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "show_doc(AutoModel)"
    ]
   },
   {
@@ -162,40 +311,41 @@
    "source": [
     "#| export\n",
     "class AutoMLForecast:\n",
-    "\n",
+    "    \"\"\"Hyperparameter optimization helper\n",
+    "    \n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    models : list or dict\n",
+    "        Auto models to be optimized.\n",
+    "    freq : str or int\n",
+    "        pandas' or polars' offset alias or integer denoting the frequency of the series.\n",
+    "    season_length : int\n",
+    "        Length of the seasonal period. This is used for producing the feature space.\n",
+    "    init_config : callable, optional (default=None)\n",
+    "        Function that takes an optuna trial and produces a configuration passed to the MLForecast constructor.\n",
+    "    fit_config : callable, optional (default=None)\n",
+    "        Function that takes an optuna trial and produces a configuration passed to the MLForecast fit method.\n",
+    "    num_threads : int (default=1)\n",
+    "        Number of threads to use when computing the features.\n",
+    "    \"\"\"\n",
     "    def __init__(\n",
     "        self,\n",
-    "        models_with_configs: Union[List[_ModelWithConfig], Dict[str, _ModelWithConfig]],\n",
+    "        models: Union[List[AutoModel], Dict[str, AutoModel]],\n",
     "        freq: Freq,\n",
     "        season_length: int,\n",
-    "        init_config: Optional[Callable] = None,\n",
-    "        fit_config: Optional[Callable] = None,\n",
+    "        init_config: Optional[_TrialToConfig] = None,\n",
+    "        fit_config: Optional[_TrialToConfig] = None,\n",
     "        num_threads: int = 1,\n",
     "    ):\n",
     "        self.freq = freq\n",
     "        self.season_length = season_length\n",
     "        self.num_threads = num_threads\n",
-    "        if isinstance(models_with_configs, list):\n",
-    "            names = _name_models([_get_model_name(m) for m, _ in models_with_configs])\n",
-    "            self.models_with_configs = {\n",
-    "                name: (model, config)\n",
-    "                for name, (model, config) in zip(names, models_with_configs)\n",
-    "            }\n",
-    "        elif isinstance(models_with_configs, dict):\n",
-    "            self.models_with_configs = models_with_configs\n",
+    "        if isinstance(models, list):\n",
+    "            model_names = _name_models([_get_model_name(m) for m in models])\n",
+    "            models_with_names = dict(zip(model_names, models))\n",
     "        else:\n",
-    "            raise ValueError(\n",
-    "                '`models_with_configs` should be a list of tuples '\n",
-    "                'or a dict from str to tuple.'\n",
-    "            )\n",
-    "        for name, (model, config) in self.models_with_configs.items():\n",
-    "            if config is not None:\n",
-    "                continue\n",
-    "            if name not in _default_model_spaces:\n",
-    "                raise NotImplementedError(\n",
-    "                    f\"{name} does not have a default config. Please provide one.\"\n",
-    "                )\n",
-    "            self.models_with_configs[name] = (model, _default_model_spaces[name])\n",
+    "            models_with_names = models\n",
+    "        self.models = models_with_names\n",
     "        self.init_config = init_config\n",
     "        if fit_config is not None:\n",
     "            self.fit_config = fit_config\n",
@@ -207,7 +357,7 @@
     "        h: int,\n",
     "        min_samples: int,\n",
     "        min_value: float,\n",
-    "    ) -> Callable:\n",
+    "    ) -> _TrialToConfig:\n",
     "        # target transforms  \n",
     "        candidate_targ_tfms = [\n",
     "            None,\n",
@@ -233,9 +383,7 @@
     "                )\n",
     "\n",
     "        # lags\n",
-    "        candidate_lags = [None]\n",
-    "        if self.season_length > 1:\n",
-    "            candidate_lags.append([self.season_length])\n",
+    "        candidate_lags = [None, [self.season_length]]\n",
     "        seasonality2extra_candidate_lags = {\n",
     "            7: [\n",
     "                [7, 14],\n",
@@ -246,6 +394,9 @@
     "                range(1, 25),\n",
     "                range(24, 24 * 7 + 1, 24),\n",
     "            ],\n",
+    "            52: [\n",
+    "                range(4, 53, 4),\n",
+    "            ]\n",
     "        }\n",
     "        candidate_lags.extend(\n",
     "            seasonality2extra_candidate_lags.get(self.season_length, [])\n",
@@ -293,13 +444,9 @@
     "            52: ['week', 'year'],\n",
     "            60: ['weekday', 'hour', 'second'],\n",
     "        }\n",
-    "        date_features = seasonality2date_features.get(self.season_length, None)\n",
+    "        candidate_date_features = seasonality2date_features.get(self.season_length, [])\n",
     "        if isinstance(self.freq, int):\n",
-    "            date_features = None\n",
-    "        if date_features is not None:\n",
-    "            use_date_features = trial.suggest_int('use_date_features', 0, 1)\n",
-    "            if not use_date_features:\n",
-    "                date_features = None        \n",
+    "            candidate_date_features = []\n",
     "\n",
     "        def config(trial):\n",
     "            # target transforms\n",
@@ -320,6 +467,16 @@
     "                lag_transforms = candidate_lag_tfms[lag_tfms_idx]\n",
     "            else:\n",
     "                lag_transforms = None\n",
+    "\n",
+    "            # date features\n",
+    "            if candidate_date_features:\n",
+    "                use_date_features = trial.suggest_int('use_date_features', 0, 1)\n",
+    "                if use_date_features:\n",
+    "                    date_features = candidate_date_features\n",
+    "                else:\n",
+    "                    date_features = None        \n",
+    "            else:\n",
+    "                date_features = None\n",
     "            \n",
     "            return {\n",
     "                'lags': lags,\n",
@@ -336,13 +493,45 @@
     "        n_windows: int,\n",
     "        h: int,\n",
     "        num_samples: int,        \n",
-    "        loss: Optional[Callable] = None,\n",
+    "        loss: Optional[Callable[[DataFrame, DataFrame], float]] = None,\n",
     "        id_col: str = 'unique_id',\n",
     "        time_col: str = 'ds',\n",
     "        target_col: str = 'y',\n",
     "        study_kwargs: Optional[Dict[str, Any]] = None,\n",
     "        optimize_kwargs: Optional[Dict[str, Any]] = None,\n",
     "    ) -> 'AutoMLForecast':\n",
+    "        \"\"\"Carry out the optimization process.\n",
+    "        Each model is optimized independently and the best one is trained on all data\n",
+    "        \n",
+    "        Parameters\n",
+    "        ----------\n",
+    "        df : pandas or polars DataFrame\n",
+    "            Series data in long format.\n",
+    "        n_windows : int\n",
+    "            Number of windows to evaluate.\n",
+    "        h : int\n",
+    "            Forecast horizon.\n",
+    "        num_samples : int\n",
+    "            Number of trials to run\n",
+    "        loss : callable, optional (default=None)\n",
+    "            Function that takes the validation and train dataframes and produces a float.\n",
+    "            If `None` will use the average SMAPE across series.\n",
+    "        id_col : str (default='unique_id')\n",
+    "            Column that identifies each serie.\n",
+    "        time_col : str (default='ds')\n",
+    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
+    "        target_col : str (default='y')\n",
+    "            Column that contains the target.        \n",
+    "        study_kwargs : dict, optional (default=None)\n",
+    "            Keyword arguments to be passed to the optuna.Study constructor.\n",
+    "        optimize_kwargs : dict, optional (default=None)\n",
+    "            Keyword arguments to be passed to the Study.optimize method.\n",
+    "\n",
+    "        Returns\n",
+    "        -------\n",
+    "        AutoMLForecast\n",
+    "            object with best models and optimization results\n",
+    "        \"\"\"\n",
     "        validate_freq(df[time_col], self.freq)\n",
     "        if self.init_config is not None:\n",
     "            init_config = self.init_config\n",
@@ -363,17 +552,13 @@
     "        if 'sampler' not in study_kwargs:\n",
     "            # for reproducibility\n",
     "            study_kwargs['sampler'] = optuna.samplers.TPESampler(seed=0)\n",
-    "        if optimize_kwargs is None:\n",
-    "            optimize_kwargs = {}\n",
-    "        if 'n_jobs' not in optimize_kwargs:\n",
-    "            optimize_kwargs['n_jobs'] = 1\n",
     "\n",
     "        self.results_ = []\n",
     "        self.models_ = {}\n",
-    "        for name, (model, model_config) in self.models_with_configs.items():\n",
+    "        for name, auto_model in self.models.items():\n",
     "            def config_fn(trial: optuna.Trial) -> float:\n",
     "                return {\n",
-    "                    'model_params': model_config(trial),\n",
+    "                    'model_params': auto_model.config(trial),\n",
     "                    'mlf_init_params': {\n",
     "                        **init_config(trial),\n",
     "                        'num_threads': self.num_threads,\n",
@@ -384,8 +569,8 @@
     "            objective = mlforecast_objective(\n",
     "                df=df,\n",
     "                config_fn=config_fn,\n",
-    "                eval_fn=loss,\n",
-    "                model=model,\n",
+    "                loss=loss,\n",
+    "                model=auto_model.model,\n",
     "                freq=self.freq,\n",
     "                n_windows=n_windows,\n",
     "                h=h,\n",
@@ -397,7 +582,7 @@
     "            study.optimize(objective, n_trials=num_samples, **optimize_kwargs)\n",
     "            self.results_.append(study)\n",
     "            best_config = study.best_trial.user_attrs['config']            \n",
-    "            best_model = clone(model)\n",
+    "            best_model = clone(auto_model.model)\n",
     "            best_model.set_params(**best_config['model_params'])\n",
     "            self.models_[name] = MLForecast(\n",
     "                models={name: best_model},\n",
@@ -415,6 +600,20 @@
     "        h: int,\n",
     "        X_df: Optional[DataFrame] = None,\n",
     "    ) -> DataFrame:\n",
+    "        \"\"\"\"Compute forecasts\n",
+    "\n",
+    "        Parameters\n",
+    "        ----------\n",
+    "        h : int\n",
+    "            Number of periods to predict.\n",
+    "        X_df : pandas or polars DataFrame, optional (default=None)\n",
+    "            Dataframe with the future exogenous features. Should have the id column and the time column.\n",
+    "\n",
+    "        Returns\n",
+    "        -------\n",
+    "        pandas or polars DataFrame\n",
+    "            Predictions for each serie and timestep, with one column per model.\n",
+    "        \"\"\"\n",
     "        all_preds = None\n",
     "        for name, model in self.models_.items():\n",
     "            preds = model.predict(h=h, X_df=X_df)\n",
@@ -424,11 +623,87 @@
     "                all_preds[name] = preds[name]\n",
     "        return all_preds\n",
     "\n",
-    "    def save(self, path: str) -> None:\n",
+    "    def save(self, path: Union[str, Path]) -> None:\n",
+    "        \"\"\"Save MLForecast objects\n",
+    "\n",
+    "        Parameters\n",
+    "        ----------\n",
+    "        path : str or pathlib.Path\n",
+    "            Directory where artifacts will be stored.\"\"\"\n",
     "        for name, model in self.models_.items():\n",
     "            model.save(f'{path}/{name}')"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2e311237",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "---\n",
+       "\n",
+       "[source](https://github.com/Nixtla/mlforecast/blob/main/mlforecast/auto.py#L242){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
+       "\n",
+       "### AutoMLForecast\n",
+       "\n",
+       ">      AutoMLForecast\n",
+       ">                      (models:Union[List[__main__.AutoModel],Dict[str,__main__.\n",
+       ">                      AutoModel]], freq:Union[int,str], season_length:int, init\n",
+       ">                      _config:Optional[Callable[[optuna.trial._trial.Trial],Dic\n",
+       ">                      t[str,Any]]]=None, fit_config:Optional[Callable[[optuna.t\n",
+       ">                      rial._trial.Trial],Dict[str,Any]]]=None,\n",
+       ">                      num_threads:int=1)\n",
+       "\n",
+       "Hyperparameter optimization helper\n",
+       "\n",
+       "|    | **Type** | **Default** | **Details** |\n",
+       "| -- | -------- | ----------- | ----------- |\n",
+       "| models | Union |  | Auto models to be optimized. |\n",
+       "| freq | Union |  | pandas' or polars' offset alias or integer denoting the frequency of the series. |\n",
+       "| season_length | int |  | Length of the seasonal period. This is used for producing the feature space. |\n",
+       "| init_config | Optional | None | Function that takes an optuna trial and produces a configuration passed to the MLForecast constructor. |\n",
+       "| fit_config | Optional | None | Function that takes an optuna trial and produces a configuration passed to the MLForecast fit method. |\n",
+       "| num_threads | int | 1 | Number of threads to use when computing the features. |"
+      ],
+      "text/plain": [
+       "---\n",
+       "\n",
+       "[source](https://github.com/Nixtla/mlforecast/blob/main/mlforecast/auto.py#L242){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
+       "\n",
+       "### AutoMLForecast\n",
+       "\n",
+       ">      AutoMLForecast\n",
+       ">                      (models:Union[List[__main__.AutoModel],Dict[str,__main__.\n",
+       ">                      AutoModel]], freq:Union[int,str], season_length:int, init\n",
+       ">                      _config:Optional[Callable[[optuna.trial._trial.Trial],Dic\n",
+       ">                      t[str,Any]]]=None, fit_config:Optional[Callable[[optuna.t\n",
+       ">                      rial._trial.Trial],Dict[str,Any]]]=None,\n",
+       ">                      num_threads:int=1)\n",
+       "\n",
+       "Hyperparameter optimization helper\n",
+       "\n",
+       "|    | **Type** | **Default** | **Details** |\n",
+       "| -- | -------- | ----------- | ----------- |\n",
+       "| models | Union |  | Auto models to be optimized. |\n",
+       "| freq | Union |  | pandas' or polars' offset alias or integer denoting the frequency of the series. |\n",
+       "| season_length | int |  | Length of the seasonal period. This is used for producing the feature space. |\n",
+       "| init_config | Optional | None | Function that takes an optuna trial and produces a configuration passed to the MLForecast constructor. |\n",
+       "| fit_config | Optional | None | Function that takes an optuna trial and produces a configuration passed to the MLForecast fit method. |\n",
+       "| num_threads | int | 1 | Number of threads to use when computing the features. |"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "show_doc(AutoMLForecast)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -436,9 +711,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import warnings\n",
-    "\n",
-    "import lightgbm as lgb\n",
     "from datasetsforecast.m4 import M4, M4Evaluation, M4Info\n",
     "from sklearn.linear_model import Ridge\n",
     "from sklearn.compose import ColumnTransformer\n",
@@ -465,73 +737,185 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f131da88-6243-4b26-92d0-c0a1f559234b",
+   "id": "7888cb6d",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pickle"
+    "ridge_pipeline = make_pipeline(\n",
+    "    ColumnTransformer(\n",
+    "        [('encoder', OneHotEncoder(), ['unique_id'])],\n",
+    "        remainder='passthrough',\n",
+    "    ),\n",
+    "    Ridge()\n",
+    ")\n",
+    "auto_ridge = AutoModel(ridge_pipeline, lambda trial: {f'ridge__{k}': v for k, v in ridge_space(trial).items()})"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "49ccb9b1-484a-4938-a762-1a843a716983",
+   "id": "88b2aa47",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>unique_id</th>\n",
+       "      <th>ds</th>\n",
+       "      <th>lgb</th>\n",
+       "      <th>ridge</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>W1</td>\n",
+       "      <td>2180</td>\n",
+       "      <td>35529.435224</td>\n",
+       "      <td>36111.103819</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>W1</td>\n",
+       "      <td>2181</td>\n",
+       "      <td>35521.764894</td>\n",
+       "      <td>36195.480607</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>W1</td>\n",
+       "      <td>2182</td>\n",
+       "      <td>35537.417268</td>\n",
+       "      <td>36107.924714</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>W1</td>\n",
+       "      <td>2183</td>\n",
+       "      <td>35538.058206</td>\n",
+       "      <td>36027.610599</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>W1</td>\n",
+       "      <td>2184</td>\n",
+       "      <td>35614.611211</td>\n",
+       "      <td>36093.419431</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4662</th>\n",
+       "      <td>W99</td>\n",
+       "      <td>2292</td>\n",
+       "      <td>15071.536978</td>\n",
+       "      <td>15319.461233</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4663</th>\n",
+       "      <td>W99</td>\n",
+       "      <td>2293</td>\n",
+       "      <td>15058.145278</td>\n",
+       "      <td>15299.892278</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4664</th>\n",
+       "      <td>W99</td>\n",
+       "      <td>2294</td>\n",
+       "      <td>15042.493434</td>\n",
+       "      <td>15272.111331</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4665</th>\n",
+       "      <td>W99</td>\n",
+       "      <td>2295</td>\n",
+       "      <td>15042.144846</td>\n",
+       "      <td>15250.457610</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4666</th>\n",
+       "      <td>W99</td>\n",
+       "      <td>2296</td>\n",
+       "      <td>15038.729044</td>\n",
+       "      <td>15232.536072</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>4667 rows × 4 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "     unique_id    ds           lgb         ridge\n",
+       "0           W1  2180  35529.435224  36111.103819\n",
+       "1           W1  2181  35521.764894  36195.480607\n",
+       "2           W1  2182  35537.417268  36107.924714\n",
+       "3           W1  2183  35538.058206  36027.610599\n",
+       "4           W1  2184  35614.611211  36093.419431\n",
+       "...        ...   ...           ...           ...\n",
+       "4662       W99  2292  15071.536978  15319.461233\n",
+       "4663       W99  2293  15058.145278  15299.892278\n",
+       "4664       W99  2294  15042.493434  15272.111331\n",
+       "4665       W99  2295  15042.144846  15250.457610\n",
+       "4666       W99  2296  15038.729044  15232.536072\n",
+       "\n",
+       "[4667 rows x 4 columns]"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "def run_group(group):\n",
-    "    train, valid = train_valid_split(group)\n",
-    "    train['unique_id'] = train['unique_id'].astype('category')\n",
-    "    valid['unique_id'] = valid['unique_id'].astype(train['unique_id'].dtype)\n",
-    "    info = M4Info[group]\n",
-    "    h = info.horizon\n",
-    "    season_length = info.seasonality\n",
-    "    ridge_pipeline = make_pipeline(\n",
-    "        ColumnTransformer(\n",
-    "            [('encoder', OneHotEncoder(), ['unique_id'])],\n",
-    "            remainder='passthrough',\n",
-    "        ),\n",
-    "        Ridge(),\n",
-    "    )\n",
-    "    auto_mlf = AutoMLForecast(\n",
-    "        freq=1,\n",
-    "        season_length=season_length,\n",
-    "        models_with_configs=[(lgb.LGBMRegressor(), None)],\n",
-    "        fit_config=lambda trial: {'static_features': ['unique_id']} if trial.suggest_int('static_features', 0, 1) else {},\n",
-    "    )\n",
-    "    with warnings.catch_warnings(record=False):\n",
-    "        warnings.simplefilter('ignore', category=UserWarning)\n",
-    "        auto_mlf.fit(\n",
-    "            df=train,\n",
-    "            n_windows=2,\n",
-    "            h=h,\n",
-    "            num_samples=30,\n",
-    "            optimize_kwargs={'timeout': 60 * 60},\n",
-    "        )\n",
-    "    with open(f'{group}_opt.pkl', 'wb') as f:\n",
-    "        pickle.dump(auto_mlf.results_, f)\n",
-    "    preds = auto_mlf.predict(h=h)\n",
-    "    for model in auto_mlf.models_with_configs.keys():\n",
-    "        print(model)\n",
-    "        print(M4Evaluation.evaluate('data', group, preds[model].values.reshape(-1, h)))    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e02af992-d4a0-492f-b929-cd9309687daf",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import time\n",
-    "\n",
     "optuna.logging.set_verbosity(optuna.logging.ERROR)\n",
-    "\n",
-    "for group in ('Yearly', 'Quarterly', 'Weekly', 'Hourly', 'Monthly'):\n",
-    "    print(f'Running {group}')\n",
-    "    start = time.perf_counter()\n",
-    "    run_group(group)\n",
-    "    print(f'{group} took {(time.perf_counter() - start) / 60:.1f} minutes')"
+    "group = 'Weekly'\n",
+    "train, valid = train_valid_split(group)\n",
+    "train['unique_id'] = train['unique_id'].astype('category')\n",
+    "valid['unique_id'] = valid['unique_id'].astype(train['unique_id'].dtype)\n",
+    "info = M4Info[group]\n",
+    "h = info.horizon\n",
+    "season_length = info.seasonality\n",
+    "auto_mlf = AutoMLForecast(\n",
+    "    freq=1,\n",
+    "    season_length=season_length,\n",
+    "    models={\n",
+    "        'lgb': AutoLightGBM(),\n",
+    "        'ridge': auto_ridge,\n",
+    "    },\n",
+    "    fit_config=lambda trial: {'static_features': ['unique_id']},\n",
+    "    num_threads=2,\n",
+    ")\n",
+    "auto_mlf.fit(\n",
+    "    df=train,\n",
+    "    n_windows=2,\n",
+    "    h=h,\n",
+    "    num_samples=2,\n",
+    "    optimize_kwargs={'timeout': 60},\n",
+    ")\n",
+    "auto_mlf.predict(h)"
    ]
   }
  ],
diff --git a/nbs/compat.ipynb b/nbs/compat.ipynb
index d848bcb6..3ad79a59 100644
--- a/nbs/compat.ipynb
+++ b/nbs/compat.ipynb
@@ -20,12 +20,36 @@
     "#| export\n",
     "try:\n",
     "    from catboost import CatBoostRegressor\n",
+    "except ImportError:\n",
+    "    class CatBoostRegressor:\n",
+    "        def __init__(self, *args, **kwargs):\n",
+    "            raise ImportError(\n",
+    "                \"Please install catboost to use this model.\"\n",
+    "            )\n",
+    "\n",
+    "try:\n",
+    "    from lightgbm import LGBMRegressor\n",
+    "except ImportError:\n",
+    "    class LGBMRegressor:\n",
+    "        def __init__(self, *args, **kwargs):\n",
+    "            raise ImportError(\n",
+    "                \"Please install lightgbm to use this model.\"\n",
+    "            )\n",
+    "\n",
+    "try:\n",
     "    from window_ops.shift import shift_array\n",
     "except ImportError:\n",
     "    def shift_array(*_args, **_kwargs):\n",
     "        raise Exception\n",
     "\n",
-    "    class CatBoostRegressor: ..."
+    "try:\n",
+    "    from xgboost import XGBRegressor\n",
+    "except ImportError:\n",
+    "    class XGBRegressor:\n",
+    "        def __init__(self, *args, **kwargs):\n",
+    "            raise ImportError(\n",
+    "                \"Please install xgboost to use this model.\"\n",
+    "            )"
    ]
   }
  ],
diff --git a/nbs/optimization.ipynb b/nbs/optimization.ipynb
index dd3422d3..68b217b6 100644
--- a/nbs/optimization.ipynb
+++ b/nbs/optimization.ipynb
@@ -31,20 +31,30 @@
    "source": [
     "#| export\n",
     "import copy\n",
-    "from typing import Callable, List, Optional\n",
+    "from typing import Any, Callable, Dict, Optional\n",
     "\n",
     "import numpy as np\n",
     "import optuna\n",
     "import utilsforecast.processing as ufp\n",
     "from sklearn.base import BaseEstimator, clone\n",
     "from utilsforecast.compat import DataFrame\n",
-    "from utilsforecast.losses import smape\n",
     "\n",
     "from mlforecast import MLForecast\n",
     "from mlforecast.compat import CatBoostRegressor\n",
     "from mlforecast.core import Freq"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f7e0e555",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| exporti\n",
+    "_TrialToConfig = Callable[[optuna.Trial], Dict[str, Any]]"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -55,8 +65,8 @@
     "#| export\n",
     "def mlforecast_objective(\n",
     "    df: DataFrame,\n",
-    "    config_fn: Callable,\n",
-    "    eval_fn: Callable,    \n",
+    "    config_fn: _TrialToConfig,\n",
+    "    loss: Callable[[DataFrame, DataFrame], float],\n",
     "    model: BaseEstimator,\n",
     "    freq: Freq,\n",
     "    n_windows: int,\n",
@@ -64,7 +74,34 @@
     "    id_col: str = 'unique_id',\n",
     "    time_col: str = 'ds',\n",
     "    target_col: str = 'y',\n",
-    ") -> Callable:\n",
+    ") -> _TrialToConfig:\n",
+    "    \"\"\"optuna objective function for the MLForecast class\n",
+    "    \n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    config_fn : callable\n",
+    "        Function that takes an optuna trial and produces a configuration with the following keys:\n",
+    "        - model_params\n",
+    "        - mlf_init_params\n",
+    "        - mlf_fit_params\n",
+    "    loss : callable\n",
+    "        Function that takes the validation and train dataframes and produces a float.\n",
+    "    model : BaseEstimator\n",
+    "        scikit-learn compatible model to be trained\n",
+    "    freq : str or int\n",
+    "        pandas' or polars' offset alias or integer denoting the frequency of the series.\n",
+    "    n_windows : int\n",
+    "        Number of windows to evaluate.\n",
+    "    h : int\n",
+    "        Forecast horizon. \n",
+    "    id_col : str (default='unique_id')\n",
+    "        Column that identifies each serie.\n",
+    "    time_col : str (default='ds')\n",
+    "        Column that identifies each timestep, its values can be timestamps or integers.\n",
+    "    target_col : str (default='y')\n",
+    "        Column that contains the target.        \n",
+    "    study_kwargs : dict, optional (default=None)\n",
+    "    \"\"\"\n",
     "    def objective(trial: optuna.Trial) -> float:\n",
     "        config = config_fn(trial)\n",
     "        trial.set_user_attr('config', copy.deepcopy(config))\n",
@@ -130,7 +167,7 @@
     "                    \"Please verify that the passed frequency (freq) matches your series' \"\n",
     "                    \"and that there aren't any missing periods.\"                    \n",
     "                )\n",
-    "            metric = eval_fn(result, train_df=train)\n",
+    "            metric = loss(result, train_df=train)\n",
     "            metrics.append(metric)\n",
     "            trial.report(metric, step=i)\n",
     "            if trial.should_prune():\n",
@@ -139,6 +176,105 @@
     "    return objective"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5ca09079",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| hide\n",
+    "from nbdev import show_doc"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ba6c74a3",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "---\n",
+       "\n",
+       "[source](https://github.com/Nixtla/mlforecast/blob/main/mlforecast/optimization.py#L24){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
+       "\n",
+       "### mlforecast_objective\n",
+       "\n",
+       ">      mlforecast_objective\n",
+       ">                            (df:Union[pandas.core.frame.DataFrame,polars.datafr\n",
+       ">                            ame.frame.DataFrame], config_fn:Callable[[optuna.tr\n",
+       ">                            ial._trial.Trial],Dict[str,Any]], loss:Callable[[Un\n",
+       ">                            ion[pandas.core.frame.DataFrame,polars.dataframe.fr\n",
+       ">                            ame.DataFrame],Union[pandas.core.frame.DataFrame,po\n",
+       ">                            lars.dataframe.frame.DataFrame]],float],\n",
+       ">                            model:sklearn.base.BaseEstimator,\n",
+       ">                            freq:Union[int,str], n_windows:int, h:int,\n",
+       ">                            id_col:str='unique_id', time_col:str='ds',\n",
+       ">                            target_col:str='y')\n",
+       "\n",
+       "optuna objective function for the MLForecast class\n",
+       "\n",
+       "|    | **Type** | **Default** | **Details** |\n",
+       "| -- | -------- | ----------- | ----------- |\n",
+       "| df | Union |  |  |\n",
+       "| config_fn | Callable |  | Function that takes an optuna trial and produces a configuration with the following keys:<br>- model_params<br>- mlf_init_params<br>- mlf_fit_params |\n",
+       "| loss | Callable |  | Function that takes the validation and train dataframes and produces a float. |\n",
+       "| model | BaseEstimator |  | scikit-learn compatible model to be trained |\n",
+       "| freq | Union |  | pandas' or polars' offset alias or integer denoting the frequency of the series. |\n",
+       "| n_windows | int |  | Number of windows to evaluate. |\n",
+       "| h | int |  | Forecast horizon.  |\n",
+       "| id_col | str | unique_id | Column that identifies each serie. |\n",
+       "| time_col | str | ds | Column that identifies each timestep, its values can be timestamps or integers. |\n",
+       "| target_col | str | y | Column that contains the target.         |\n",
+       "| **Returns** | **Callable** |  |  |"
+      ],
+      "text/plain": [
+       "---\n",
+       "\n",
+       "[source](https://github.com/Nixtla/mlforecast/blob/main/mlforecast/optimization.py#L24){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
+       "\n",
+       "### mlforecast_objective\n",
+       "\n",
+       ">      mlforecast_objective\n",
+       ">                            (df:Union[pandas.core.frame.DataFrame,polars.datafr\n",
+       ">                            ame.frame.DataFrame], config_fn:Callable[[optuna.tr\n",
+       ">                            ial._trial.Trial],Dict[str,Any]], loss:Callable[[Un\n",
+       ">                            ion[pandas.core.frame.DataFrame,polars.dataframe.fr\n",
+       ">                            ame.DataFrame],Union[pandas.core.frame.DataFrame,po\n",
+       ">                            lars.dataframe.frame.DataFrame]],float],\n",
+       ">                            model:sklearn.base.BaseEstimator,\n",
+       ">                            freq:Union[int,str], n_windows:int, h:int,\n",
+       ">                            id_col:str='unique_id', time_col:str='ds',\n",
+       ">                            target_col:str='y')\n",
+       "\n",
+       "optuna objective function for the MLForecast class\n",
+       "\n",
+       "|    | **Type** | **Default** | **Details** |\n",
+       "| -- | -------- | ----------- | ----------- |\n",
+       "| df | Union |  |  |\n",
+       "| config_fn | Callable |  | Function that takes an optuna trial and produces a configuration with the following keys:<br>- model_params<br>- mlf_init_params<br>- mlf_fit_params |\n",
+       "| loss | Callable |  | Function that takes the validation and train dataframes and produces a float. |\n",
+       "| model | BaseEstimator |  | scikit-learn compatible model to be trained |\n",
+       "| freq | Union |  | pandas' or polars' offset alias or integer denoting the frequency of the series. |\n",
+       "| n_windows | int |  | Number of windows to evaluate. |\n",
+       "| h | int |  | Forecast horizon.  |\n",
+       "| id_col | str | unique_id | Column that identifies each serie. |\n",
+       "| time_col | str | ds | Column that identifies each timestep, its values can be timestamps or integers. |\n",
+       "| target_col | str | y | Column that contains the target.         |\n",
+       "| **Returns** | **Callable** |  |  |"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "show_doc(mlforecast_objective)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -249,7 +385,7 @@
     "        }\n",
     "    }\n",
     "\n",
-    "def eval_fn(df, _train_df=None):\n",
+    "def loss(df, train_df):\n",
     "    return smape(df, models=['model'])['model'].mean()"
    ]
   },
@@ -258,29 +394,6 @@
    "execution_count": null,
    "id": "23296223-7aad-4333-a43f-0a22d51f1163",
    "metadata": {},
-   "outputs": [],
-   "source": [
-    "optuna.logging.set_verbosity(optuna.logging.WARNING)\n",
-    "objective = mlforecast_objective(\n",
-    "    df=weekly_train,\n",
-    "    config_fn=config_fn,\n",
-    "    eval_fn=eval_fn,    \n",
-    "    model=lgb.LGBMRegressor(),\n",
-    "    freq=1,\n",
-    "    n_windows=2,\n",
-    "    h=h,\n",
-    ")\n",
-    "study = optuna.create_study(\n",
-    "    direction='minimize', sampler=optuna.samplers.TPESampler(seed=0)\n",
-    ")\n",
-    "study.optimize(objective, n_trials=10, n_jobs=1)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "26bed3c1-f819-4e73-a0de-dbd3ebbd0fc3",
-   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -311,9 +424,9 @@
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>Weekly</th>\n",
-       "      <td>7.880505</td>\n",
-       "      <td>2.181829</td>\n",
-       "      <td>0.822896</td>\n",
+       "      <td>9.261538</td>\n",
+       "      <td>2.614473</td>\n",
+       "      <td>0.976158</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -321,7 +434,7 @@
       ],
       "text/plain": [
        "           SMAPE      MASE       OWA\n",
-       "Weekly  7.880505  2.181829  0.822896"
+       "Weekly  9.261538  2.614473  0.976158"
       ]
      },
      "execution_count": null,
@@ -330,6 +443,20 @@
     }
    ],
    "source": [
+    "optuna.logging.set_verbosity(optuna.logging.WARNING)\n",
+    "objective = mlforecast_objective(\n",
+    "    df=weekly_train,\n",
+    "    config_fn=config_fn,\n",
+    "    loss=loss,    \n",
+    "    model=lgb.LGBMRegressor(),\n",
+    "    freq=1,\n",
+    "    n_windows=2,\n",
+    "    h=h,\n",
+    ")\n",
+    "study = optuna.create_study(\n",
+    "    direction='minimize', sampler=optuna.samplers.TPESampler(seed=0)\n",
+    ")\n",
+    "study.optimize(objective, n_trials=2)\n",
     "best_cfg = study.best_trial.user_attrs['config']\n",
     "final_model = MLForecast(\n",
     "    models=[lgb.LGBMRegressor(**best_cfg['model_params'])],\n",