diff --git a/examples/examples_uq/plot_nas_deep_ensemble_uq_regression_pytorch_error_model.py b/examples/examples_uq/plot_nas_deep_ensemble_uq_regression_pytorch_error_model.py
new file mode 100644
index 00000000..627a778c
--- /dev/null
+++ b/examples/examples_uq/plot_nas_deep_ensemble_uq_regression_pytorch_error_model.py
@@ -0,0 +1,1333 @@
+r"""
+Neural Architecture Search and Deep Ensemble with Uncertainty Quantification for Regression (Pytorch)
+=====================================================================================================
+
+**Author(s)**: Romain Egele, Brett Eiffert.
+
+In this tutorial, you will learn how to perform **Neural Architecture Search (NAS)** and use it to construct a diverse deep ensemble with disentangled **aleatoric** and **epistemic uncertainty**.
+ 
+NAS is the idea of automatically optimizing the architecture of deep neural networks to solve a given task. Here, we will use **hyperparameter optimization (HPO)** algorithms to guide the NAS process.
+
+Specifically, in this tutorial you will learn how to:
+
+1.	**Define a customizable PyTorch module** that exposes neural architecture hyperparameters.
+2.	**Define constraints** on the neural architecture hyperparameters to reduce redundancies and improve efficiency of the optimization.
+
+This tutorial will provide a hands-on approach to leveraging NAS for robust regression models with well-calibrated uncertainty estimates.
+
+"""
+# %%
+# Installation and imports
+# ------------------------
+#
+# Installing dependencies with the :ref:`pip installation <install-pip>` is recommended. It requires **Python >= 3.10**.
+#
+# .. code-block:: bash
+#
+#     %%bash
+#     pip install "deephyper[ray,torch]"
+
+#%%
+
+# .. dropdown:: Import statements
+import json
+import os
+import pathlib
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+from sklearn.model_selection import train_test_split
+from tqdm import tqdm
+
+WIDTH_PLOTS = 8
+HEIGHT_PLOTS = WIDTH_PLOTS / 1.618
+
+# %%
+# Synthetic data generation
+# -------------------------
+# 
+# We generate synthetic data from a 1D scalar function :math:`Y = f(X) + \epsilon(X)`, where :math:`X,Y` are random variables with support :math:`\mathbb{R}`.
+# 
+# The training data are drown uniformly from :math:`X \sim U([-30,-15] \cup [15,30])` with:
+# 
+# .. math::
+# 
+#     f(x) = \cos(x/2) + 2 \cdot \sin(x/10) + x/100
+# 
+# and :math:`\epsilon(X) \sim \mathcal{N}(0, \sigma(X))` with:
+#
+# - :math:`\sigma(x) = 0.5` if :math:`x \in [-30,-15]`
+# - :math:`\sigma(x) = 1.0` if :math:`x \in [15,30]`
+
+# .. dropdown:: Loading synthetic data
+def load_data(
+    developement_size=500,
+    test_size=200,
+    random_state=42,
+    x_min=-50,
+    x_max=50,
+):
+    rs = np.random.RandomState(random_state)
+
+    def f(x):
+        return np.cos(x / 2) + 2 * np.sin(x / 10) + x / 100
+
+    x_1 = rs.uniform(low=-30, high=-15.0, size=developement_size // 2)
+    eps_1 = rs.normal(loc=0.0, scale=0.5, size=developement_size // 2)
+    y_1 = f(x_1) + eps_1
+
+    x_2 = rs.uniform(low=15.0, high=30.0, size=developement_size // 2)
+    eps_2 = rs.normal(loc=0.0, scale=1.0, size=developement_size // 2)
+    y_2 = f(x_2) + eps_2
+
+    x = np.concatenate([x_1, x_2], axis=0)
+    y = np.concatenate([y_1, y_2], axis=0)
+
+    test_X = np.linspace(x_min, x_max, test_size)
+    test_y = f(test_X)
+
+    x = x.reshape(-1, 1)
+    y = y.reshape(-1, 1)
+
+    train_X, valid_X, train_y, valid_y = train_test_split(
+        x, y, test_size=0.33, random_state=random_state
+    )
+
+    test_X = test_X.reshape(-1, 1)
+    test_y = test_y.reshape(-1, 1)
+
+    return (train_X, train_y), (valid_X, valid_y), (test_X, test_y)
+
+
+(train_X, train_y), (valid_X, valid_y), (test_X, test_y) = load_data()
+
+y_mu, y_std = np.mean(train_y), np.std(train_y)
+
+x_lim, y_lim = 50, 7
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+_ = plt.scatter(train_X, train_y, s=5, label="Training")
+_ = plt.scatter(valid_X, valid_y, s=5, label="Validation")
+_ = plt.plot(test_X, test_y, linestyle="--", color="gray", label="Test")
+_ = plt.fill_between([-30, -15], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.fill_between([15, 30], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.xlim(-x_lim, x_lim)
+_ = plt.ylim(-y_lim, y_lim)
+_ = plt.legend()
+_ = plt.xlabel(r"$x$")
+_ = plt.ylabel(r"$f(x)$")
+_ = plt.grid(which="both", linestyle=":")
+
+# %%
+# Configurable neural network with uncertainty
+# --------------------------------------------
+#
+# We define a configurable Pytorch module to be able to explore:
+#
+# - the number of layers
+# - the number of units per layer
+# - the activation function per layer
+# - the dropout rate
+# - the output layer
+#
+# The output of this module will be a Gaussian distribution :math:`\mathcal{N}(\mu_\theta(x), \sigma_\theta(x))`, where :math:`\theta` represent the concatenation of the weights and the hyperparameters of our model.
+#
+# The uncertainty :math:`\sigma_\theta(x)` estimated by the network is an estimator of :math:`V_Y[Y|X=x]` therefore corresponding
+# to aleatoric uncertainty (a.k.a., intrinsic noise).
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+
+
+class DeepNormalRegressor(nn.Module):
+    def __init__(
+        self,
+        n_inputs,
+        layers,
+        n_units_mean=64,
+        n_units_std=64,
+        std_offset=1e-3,
+        softplus_factor=0.05,
+        loc=0,
+        scale=1.0,
+    ):
+        super().__init__()
+
+        layers_ = []
+        prev_n_units = n_inputs
+        for n_units, activation, dropout_rate in layers:
+            linear_layer = nn.Linear(prev_n_units, n_units)
+            if activation == "relu":
+                activation_layer = nn.ReLU()
+            elif activation == "sigmoid":
+                activation_layer = nn.Sigmoid()
+            elif activation == "tanh":
+                activation_layer = nn.Tanh()
+            elif activation == "swish":
+                activation_layer = nn.SiLU()
+            elif activation == "mish":
+                activation_layer = nn.Mish()
+            elif activation == "gelu":
+                activation_layer = nn.GELU()
+            elif activation == "silu":
+                activation_layer = nn.SiLU()
+            dropout_layer = nn.Dropout(dropout_rate)
+
+            layers_.extend([linear_layer, activation_layer, dropout_layer])
+
+            prev_n_units = n_units
+
+        # Shared parameters
+        self.shared_layer = nn.Sequential(
+            *layers_,
+        )
+
+        # Mean parameters
+        self.mean_layer = nn.Sequential(
+            nn.Linear(prev_n_units, n_units_mean),
+            nn.ReLU(),
+            nn.Linear(n_units_mean, 1),
+        )
+
+        # Standard deviation parameters
+        self.std_layer = nn.Sequential(
+            nn.Linear(prev_n_units, n_units_std),
+            nn.ReLU(),
+            nn.Linear(n_units_std, 1),
+            nn.Softplus(beta=1.0, threshold=20.0),  # enforces positivity
+        )
+
+        self.std_offset = std_offset
+        self.softplus_factor = softplus_factor
+        self.loc = loc
+        self.scale = scale
+
+    def forward(self, x):
+        # Shared embedding
+        shared = self.shared_layer(x)
+
+        # Parametrization of the mean
+        mu = self.mean_layer(shared) + self.loc
+
+        # Parametrization of the standard deviation
+        sigma = self.std_offset + self.std_layer(self.softplus_factor * shared) * self.scale
+
+        return torch.distributions.Normal(mu, sigma)
+
+# %%
+# Hyperparameter search space
+# ---------------------------
+#
+# We define the hyperparameter space that includes both **neural architecture** and **training hyperparameters**.
+#
+# Without having a good heuristic on training hyperparameters given the neural architecture hyperparameter search space 
+# it is important to define them jointly with the neural architecture hyperparameters as they can have strong interactions. 
+#
+# In the definition of the hyperparameter space, we add constraints using :class:`ConfigSpace.GreaterThanCondition` to
+# represent when an hyperparameter is active. In this example, "active" means it actually influence the code execution of
+# the trained model.
+
+from ConfigSpace import GreaterThanCondition
+from deephyper.hpo import HpProblem
+
+
+def create_hpo_problem(min_num_layers=3, max_num_layers=8, max_num_units=512):
+    problem = HpProblem()
+
+    # Neural Architecture Hyperparameters
+    num_layers = problem.add_hyperparameter((min_num_layers, max_num_layers), "num_layers", default_value=5)
+
+    conditions = []
+    for i in range(max_num_layers):
+
+        # Adding the hyperparameters that impact each layer of the model
+        layer_i_units = problem.add_hyperparameter((16, max_num_units), f"layer_{i}_units", default_value=max_num_units)
+        layer_i_activation = problem.add_hyperparameter(
+            ["relu", "sigmoid", "tanh", "swish", "mish", "gelu", "silu"],
+            f"layer_{i}_activation",
+            default_value="relu",
+        )
+        layer_i_dropout_rate = problem.add_hyperparameter(
+            (0.0, 0.25), f"layer_{i}_dropout_rate", default_value=0.0
+        )
+
+        # Adding the constraints to define when these hyperparameters are active
+        if i + 1 > min_num_layers:
+            conditions.extend(
+                [
+                    GreaterThanCondition(layer_i_units, num_layers, i),
+                    GreaterThanCondition(layer_i_activation, num_layers, i),
+                    GreaterThanCondition(layer_i_dropout_rate, num_layers, i),
+                ]
+            )
+
+    problem.add_conditions(conditions)
+
+    # Hyperparameters of the output layers
+    problem.add_hyperparameter((16, max_num_units), "n_units_mean", default_value=max_num_units)
+    problem.add_hyperparameter((16, max_num_units), "n_units_std", default_value=max_num_units)
+    problem.add_hyperparameter((1e-8, 1e-2, "log-uniform"), "std_offset", default_value=1e-3)
+    problem.add_hyperparameter((0.01, 1.0), "softplus_factor", default_value=0.05)
+
+    # Training Hyperparameters
+    problem.add_hyperparameter((1e-5, 1e-1, "log-uniform"), "learning_rate", default_value=2e-3)
+    problem.add_hyperparameter((8, 256, "log-uniform"), "batch_size", default_value=32)
+    problem.add_hyperparameter((0.01, 0.99), "lr_scheduler_factor", default_value=0.1)
+    problem.add_hyperparameter((10, 100), "lr_scheduler_patience", default_value=20)
+
+    return problem
+
+problem = create_hpo_problem()
+problem
+
+# %%
+# Loss and Metric
+# ---------------
+#
+# For the loss we will use the Gaussian negative log-likelihood to evalute the quality of the 
+# predicted distribution :math:`\mathcal{N}(\mu_\theta(x), \sigma_\theta(x))` using with formula:
+#
+# .. math::
+#
+#     L_\text{NLL}(x, y;\theta) = \frac{1}{2}\left(\log\left(\sigma_\theta^{2}(x)\right) + \frac{\left(y-\mu_{\theta}(x)\right)^{2}}{\sigma_{\theta}^{2}(x)}\right) + \text{cst}
+# 
+# As complementary metric, we use the squared error to evaluate the quality of the mean predictions :math:`\mu_\theta(x)`:
+#
+# .. math::
+#
+#     L_\text{SE}(x, y;\theta) = (\mu_\theta(x)-y)^2
+#
+def nll(y, rv_y):
+    """Negative log likelihood for Pytorch distribution.
+
+    Args:
+        y: true data.
+        rv_y: learned (predicted) probability distribution.
+    """
+    return -rv_y.log_prob(y)
+
+
+def squared_error(y_true, rv_y):
+    """Squared error for Pytorch distribution.
+
+    Args:
+        y: true data.
+        rv_y: learned (predicted) probability distribution.
+    """
+    y_pred = rv_y.mean
+    return (y_true - y_pred) ** 2
+
+# %%
+# Training loop
+# -------------
+#
+# In our training loop, we make sure to collect training and validation learning curves for better analysis.
+#
+# We also add a mechanism to checkpoint weights of the model based on the best observed validation loss.
+#
+# Finally, we add an early stopping mechanism to save computing resources.
+
+# .. dropdown:: Training loop
+def train_one_step(model, optimizer, x_batch, y_batch):
+    model.train()
+    optimizer.zero_grad()
+    y_dist = model(x_batch)
+
+    loss = torch.mean(nll(y_batch, y_dist))
+    mse = torch.mean(squared_error(y_batch, y_dist))
+
+    loss.backward()
+    optimizer.step()
+
+    return loss, mse
+
+
+def train(
+    job,
+    model,
+    optimizer,
+    x_train,
+    x_val,
+    y_train,
+    y_val,
+    n_epochs,
+    batch_size,
+    scheduler=None,
+    patience=200,
+    progressbar=True,
+):
+    data_train = DataLoader(TensorDataset(x_train, y_train), batch_size=batch_size, shuffle=True)
+
+    checkpointed_state_dict = model.state_dict()
+    checkpointed_val_loss = np.inf
+
+    train_loss, val_loss = [], []
+    train_mse, val_mse = [], []
+
+    tqdm_bar = tqdm(total=n_epochs, disable=not progressbar)
+
+    for epoch in range(n_epochs):
+        batch_losses_t, batch_losses_v, batch_mse_t, batch_mse_v = [], [], [], []
+
+        for batch_x, batch_y in data_train:
+            b_train_loss, b_train_mse = train_one_step(model, optimizer, batch_x, batch_y)
+
+            model.eval()
+            y_dist = model(x_val)
+            b_val_loss = torch.mean(nll(y_val, y_dist))
+            b_val_mse = torch.mean(squared_error(y_val, y_dist))
+
+            batch_losses_t.append(to_numpy(b_train_loss))
+            batch_mse_t.append(to_numpy(b_train_mse))
+            batch_losses_v.append(to_numpy(b_val_loss))
+            batch_mse_v.append(to_numpy(b_val_mse))
+
+        train_loss.append(np.mean(batch_losses_t))
+        val_loss.append(np.mean(batch_losses_v))
+        train_mse.append(np.mean(batch_mse_t))
+        val_mse.append(np.mean(batch_mse_v))
+
+        if scheduler is not None:
+            scheduler.step(val_loss[-1])
+
+        tqdm_bar.update(1)
+        tqdm_bar.set_postfix(
+            {
+                "train_loss": f"{train_loss[-1]:.3f}",
+                "val_loss": f"{val_loss[-1]:.3f}",
+                "train_mse": f"{train_mse[-1]:.3f}",
+                "val_mse": f"{val_mse[-1]:.3f}",
+            }
+        )
+
+        # Checkpoint weights if they improve
+        if val_loss[-1] < checkpointed_val_loss:
+            checkpointed_val_loss = val_loss[-1]
+            checkpointed_state_dict = model.state_dict()
+
+        # Early discarding
+        job.record(budget=epoch+1, objective=-val_loss[-1])
+        if job.stopped():
+            break
+
+        if len(val_loss) > (patience + 1) and val_loss[-patience - 1] < min(val_loss[-patience:]):
+            break
+
+    # Reload the best weights
+    model.load_state_dict(checkpointed_state_dict)
+
+    return train_loss, val_loss, train_mse, val_mse
+
+# %% 
+# Run time
+# --------
+import multiprocessing
+
+dtype = torch.float32
+if torch.cuda.is_available():
+    device = "cuda"
+    device_count = 1
+else:
+    device = "cpu"
+    device_count = multiprocessing.cpu_count()
+
+print(f"Runtime with {device=}, {device_count=}, {dtype=}")
+
+# %%
+
+# .. dropdown:: Conversion utility functions
+
+def to_torch(array):
+    return torch.from_numpy(array).to(device=device, dtype=dtype)
+
+def to_numpy(tensor):
+    return tensor.detach().cpu().numpy()
+
+# %%
+# Evaluation function
+# -------------------
+#
+# We start by defining a function that will create the Torch module from a dictionnary of hyperparameters.
+
+
+def create_model(parameters: dict, y_mu=0, y_std=1):
+    num_layers = parameters["num_layers"]
+    torch_module = DeepNormalRegressor(
+        n_inputs=1,
+        layers=[
+            (
+                parameters[f"layer_{i}_units"],
+                parameters[f"layer_{i}_activation"],
+                parameters[f"layer_{i}_dropout_rate"],
+            )
+            for i in range(num_layers)
+        ],
+        n_units_mean=parameters["n_units_mean"],
+        n_units_std=parameters["n_units_std"],
+        std_offset=parameters["std_offset"],
+        softplus_factor=parameters["softplus_factor"],
+        loc=y_mu,
+        scale=y_std,
+    ).to(device=device, dtype=dtype)
+    return torch_module
+
+# %%
+#
+# The evaluation function (often called ``run``-function in DeepHyper) is the function that 
+# receives suggested parameters as inputs ``job.parameters`` and returns an ``"objective"`` 
+# that we want to maximize.
+
+# max_n_epochs = 1_000
+max_n_epochs = 1
+
+
+def run(job, model_checkpoint_dir=".", verbose=False):
+    (x, y), (vx, vy), (tx, ty) = load_data()
+
+    # Create the model based on neural architecture hyperparameters
+    model = create_model(job.parameters, y_mu, y_std)
+
+    if verbose:
+        print(model)
+
+    # Initialize training loop based on training hyperparameters
+    optimizer = torch.optim.Adam(model.parameters(), lr=job.parameters["learning_rate"])
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer,
+        factor=job.parameters["lr_scheduler_factor"],
+        patience=job.parameters["lr_scheduler_patience"],
+    )
+
+    x, vx, tx = to_torch(x), to_torch(vx), to_torch(tx)
+    y, vy, ty = to_torch(y), to_torch(vy), to_torch(ty)
+
+    try:
+        train_losses, val_losses, train_mse, val_mse = train(
+            job,
+            model,
+            optimizer,
+            x,
+            vx,
+            y,
+            vy,
+            n_epochs=max_n_epochs,
+            batch_size=job.parameters["batch_size"],
+            scheduler=scheduler,
+            progressbar=verbose,
+        )
+    except Exception:
+        return "F_fit"
+
+    ty_pred = model(tx)
+    test_loss = to_numpy(torch.mean(nll(ty, ty_pred)))
+    test_mse = to_numpy(torch.mean(squared_error(ty, ty_pred)))
+
+    # Saving the model's state (i.e., weights)
+    torch.save(model.state_dict(), os.path.join(model_checkpoint_dir, f"model_{job.id}.pt"))
+
+    return {
+        "objective": -val_losses[-1],
+        "metadata": {
+            "train_loss": train_losses,
+            "val_loss": val_losses,
+            "train_mse": train_mse,
+            "val_mse": val_mse,
+            "test_loss": test_loss,
+            "test_mse": test_mse,
+            "budget": len(val_losses),
+        },
+    }
+
+# %% 
+# Evaluation of the baseline
+# --------------------------
+#
+# We evaluate the default configuration of hyperparameters that we call "baseline" using the same evaluation function.
+# This allows to test the evaluation function.
+
+from deephyper.evaluator import RunningJob
+
+baseline_dir = "nas_baseline_regression"
+
+def evaluate_baseline(problem):
+    model_checkpoint_dir = os.path.join(baseline_dir, "models")
+    pathlib.Path(model_checkpoint_dir).mkdir(parents=True, exist_ok=True)
+
+    default_parameters = problem.default_configuration
+    print(f"{default_parameters=}\n")
+
+    result = run(
+        RunningJob(parameters=default_parameters),
+        model_checkpoint_dir=model_checkpoint_dir,
+        verbose=True,
+    )
+    return result
+
+baseline_results = evaluate_baseline(problem)
+
+# %%
+# Then, we look at the learning curves of our baseline model returned by the evaluation function.
+#
+# These curves display a good learning behaviour:
+# 
+# - the training and validation curves follow each other closely and are decreasing.
+# - a clear convergence plateau is reached at the end of the training.
+
+# .. dropdown:: Make learning curves plot
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+
+x_values = np.arange(1, len(baseline_results["metadata"]["train_loss"]) + 1)
+_ = plt.plot(
+    x_values,
+    baseline_results["metadata"]["train_loss"],
+    label="Training",
+)
+_ = plt.plot(
+    x_values,
+    baseline_results["metadata"]["val_loss"],
+    label="Validation",
+)
+
+_ = plt.xlim(x_values.min(), x_values.max())
+_ = plt.grid(which="both", linestyle=":")
+_ = plt.legend()
+_ = plt.xlabel("Epochs")
+_ = plt.ylabel("NLL")
+
+
+# %%
+# In addition, we look at the predictions by reloading the checkpointed weights.
+#
+# We first need to recreate the torch module and then we update its state using the checkpointed weights.
+
+weights_path = os.path.join(baseline_dir, "models",  "model_0.0.pt")
+parameters = problem.default_configuration
+torch_module = create_model(parameters, y_mu, y_std)
+torch_module.load_state_dict(torch.load(weights_path, weights_only=True))
+torch_module.eval()
+
+y_pred = torch_module.forward(to_torch(test_X))
+y_pred_mean = to_numpy(y_pred.loc)
+y_pred_std = to_numpy(y_pred.scale)
+
+# %%
+
+# .. dropdown:: Make prediction plot
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+_ = plt.scatter(train_X, train_y, s=5, label="Training")
+_ = plt.scatter(valid_X, valid_y, s=5, label="Validation")
+_ = plt.plot(test_X, test_y, linestyle="--", color="gray", label="Test")
+
+_ = plt.plot(test_X, y_pred_mean, label=r"$\mu(x)$")
+kappa = 1.96
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred_mean - kappa * y_pred_std).reshape(-1),
+    (y_pred_mean + kappa * y_pred_std).reshape(-1),
+    alpha=0.25,
+    label=r"$\sigma_\text{al}(x)$",
+)
+
+_ = plt.fill_between([-30, -15], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.15)
+_ = plt.fill_between([15, 30], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.15)
+_ = plt.xlim(-x_lim, x_lim)
+_ = plt.ylim(-y_lim, y_lim)
+_ = plt.legend(ncols=2)
+_ = plt.xlabel(r"$x$")
+_ = plt.ylabel(r"$f(x)$")
+_ = plt.grid(which="both", linestyle=":")
+
+# %%
+# Neural architecture search
+# --------------------------
+#
+# We will now use Bayesian opimization to perform neural architecture search. 
+# The sequential Bayesian optimization algorithm can be described by the following pseudo-code:
+#
+# Sequential Bayesian optimization
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# **Algorithm**: Bayesian Optimization (a.k.a., Efficient Global Optimization (EGO))
+#
+#   Inputs
+#     :math:`\texttt{thetaSpace}`: a hyperparameter space
+#
+#     :math:`\texttt{nInitial}`: the number of initial hyperparameter configurations
+#
+#     :math:`\texttt{f}`: a function that returns the objective of the learning workflow
+#
+#   Outputs
+#     :math:`\texttt{thetaStar}` the recommended hyperparameter configuration
+#
+#   :math:`\texttt{thetaArray}, \texttt{objArray} \gets` New empty arrays of hyperparameter configurations and objectives
+#   :math:`\texttt{model} \gets` New surrogate model
+#
+#   Loop until stopping criteria is not valid
+#
+#     If Length of :math:`\texttt{thetaArray} < \texttt{nInitial}` then
+#
+#       :math:`\texttt{theta} \gets` Sample hyperparameter configuration from :math:`\texttt{thetaSpace}`
+#
+#     Else
+#
+#       Update :math:`\texttt{model}` with :math:`\texttt{thetaArray}, \texttt{objArray}`
+#
+#       :math:`\texttt{theta} \gets` Returns :math:`\texttt{theta}` in :math:`\texttt{thetaSpace}` that maximizes 
+#       the acquisition function for the current :math:`\texttt{model}`
+#
+#     :math:`\texttt{obj} \gets` Returns the objective of learning workflow :math:`\texttt{f}(\texttt{theta})`
+#
+#     :math:`\texttt{thetaArray}  \gets` Concatenate :math:`\texttt{thetaArray}` with :math:`[\texttt{theta}]`
+#
+#     :math:`\texttt{objArray}  \gets` Concatenate :math:`\texttt{objArray}` with :math:`[\texttt{obj}]`
+#
+#     :math:`\texttt{thetaStar} \gets` Update recommendation
+#
+# Parallel Bayesian optimization
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# In DeepHyper, instead of just performing sequential Bayesian optimization we provide asynchronous parallelisation for
+# Bayesian optimization (and other methods). This allows to execute multiple evaluation function in parallel to collect observations of objectives
+# faster.
+#
+# In this example, we will focus on using centralized Bayesian optimization (CBO). In this setting, we have one main process that runs the
+# Bayesian optimization algorithm and we have multiple worker processes that run evaluation functions. The class we use for this is
+# :class:`deephyper.hpo.CBO`.
+#
+# Let us start by explaining import configuration parameters of :class:`deephyper.hpo.CBO`:
+# 
+# - ``initial_points``: is a list of initial hyperparameter configurations to test, we add the baseline hyperparameters as we want to be at least better than this configuration.
+# - ``surrogate_model_*``: are parameters related to the surrogate model we use, here ``"ET"`` is an alias for the Extremely Randomized Trees regression model.
+# - ``multi_point_strategy``: is the strategy we use for parallel suggestion of hyperparameters, here we use the ``qUCBd`` that will sample for each new parallel configuration a different :math:`\kappa^j_i` value from an exponential with mean :math:`\kappa_i` where :math:`j` is the index in the current generated parallel batch and :math:`i` is the iteration of the Bayesian optimization loop. ``UCB`` corresponds to the Upper Confidence Bound acquisition function:
+# 
+# .. math::
+#     
+#     \alpha_\text{UCB}(\theta;\kappa) = \mu_\text{ET}(\theta) + \kappa \cdot \sigma_\text{ET}(\theta)
+#
+# where :math:`\mu_\text{ET}(\theta)` and :math:`\sigma_\text{ET}^2(\theta)` are respectively estimators of :math:`E_C[C|\Theta=\theta]` and :math:`V_C[C|\Theta=\theta]` with :math:`C` the random variable describing the objective (or cost) and :math:`\Theta` the random variable describing the hyperparameters alone.
+#
+# Finally the ``"d"`` postfix in ``qUCBd`` means that we will only consider the epistemic component of the uncertainty returned by the surrogate model.
+# Thanks to the law of total variance we have the following decomposition:
+#
+# .. math::
+# 
+#     V_C[C|\Theta=\theta] = E_\text{tree}\left[V_C[C|\Theta=\theta;\text{tree}\right] + V_\text{tree}\left[E_C[C|\Theta=\theta;\text{tree}]\right]
+# 
+# Then, we define :math:`\sigma_{\text{ET},\text{ep}}(\theta)` as the empirical estimate of :math:`V_\text{tree}\left[E_C[C|\Theta=\theta;\text{tree}]\right]`.
+# Then, we define :math:`\alpha_\text{qUCBd}(\theta;\kappa^j_i)` as:
+#
+# .. math::
+#     
+#     \alpha_\text{qUCBd}(\theta;\kappa^j_i) = \mu_\text{ET}(\theta) + \kappa^j_i \cdot \sigma_{\text{ET},\text{ep}}(\theta)
+#
+# Interestingly the same trick will be used later to decompose the uncertainty of the deep ensemble.
+#
+# - ``acq_optimizer_*``: are parameters related to optimization of the previously defined acquisition function.
+# - ``kappa`` and ``scheduler``: are the parameters that define the schedule of :math:`\kappa^j_i` previously mentionned.
+# - ``objective_scaler``: is a parameter that can be used to rescale the observed objectives (e.g., identity, min-max, log).
+search_kwargs = {
+    "initial_points": [problem.default_configuration],
+    "multi_point_strategy": "qUCBd",  # Multi-point strategy for asynchronous batch generations (explained later)
+    "acq_optimizer": "mixedga",  # Use continuous Genetic Algorithm for the acquisition function optimizer
+    "random_state": 42,  # Random seed
+}
+# %%
+# Then, we create the search instance.
+#
+# For this we pass the hyperparameter ``problem``, the ``evaluator`` and also a ``stopper`` (optional).
+#
+# The ``problem`` is the instance of :class:`deephyper.hpo.HpProblem` that we defined in previous sections.
+#
+# The ``evaluator`` is a subclass of :class:`deephyper.evaluator.Evaluator` that provides a ``.submit(...)`` method and a ``.gather(...)`` method to
+# submit and gather asynchronous evaluations.
+# 
+# The ``stopper`` is an optional parameter that allows to use an early-discarding (a.k.a., multi-fidelity) strategy to stop early low performing evaluations.
+# In our case we will use the median early-discarding strategy. 
+# This strategy consists in early stopping the training if the observed objective at the current budget is worse than the median objective for the same budget.
+from deephyper.evaluator import Evaluator
+from deephyper.evaluator.callback import TqdmCallback
+from deephyper.hpo import CBO
+from deephyper.stopper import MedianStopper
+
+
+hpo_dir = "nas_regression"
+
+
+def run_neural_architecture_search(problem, max_evals):
+    model_checkpoint_dir = os.path.join(hpo_dir, "models")
+    pathlib.Path(model_checkpoint_dir).mkdir(parents=True, exist_ok=True)
+
+    method_kwargs = {
+        "run_function_kwargs": {
+            "model_checkpoint_dir": model_checkpoint_dir,
+            "verbose": False,
+        },
+        "callbacks": [TqdmCallback()],
+    }
+
+    if device == "cuda":
+        method_kwargs.update({
+            "num_cpus": device_count,
+            "num_gpus": device_count,
+            "num_cpus_per_task": 1,
+            "num_gpus_per_task": 1,
+        })
+    else:
+        method_kwargs.update({
+            "num_cpus": device_count,
+            "num_cpus_per_task": 1,
+        })
+    
+
+    evaluator = Evaluator.create(
+        run,
+        method="ray",  
+        method_kwargs=method_kwargs,
+    )
+
+    stopper = None
+    
+    # Uncomment the following to speed-up the search
+    # stopper = MedianStopper(min_steps=50, max_steps=max_n_epochs, interval_steps=50)
+
+    search = CBO(problem, log_dir=hpo_dir, stopper=stopper, **search_kwargs)
+
+    results = search.search(evaluator, max_evals=max_evals)
+
+    return results
+
+
+# %%
+# You can download precomputed results if you want to skip the slow neural architecture search. We provide the following two set of precomputed results:
+#
+# - Link to precomputed results without stopper: ``https://drive.google.com/uc?id=1VOV-UM0ws0lopHvoYT_9RAiRdT1y4Kus``
+# - Link to precomputed results with median stopper: ``https://drive.google.com/uc?id=1I09-ZaH4BzQfBOw6YmhzgKLFWBsdrvpg``
+#
+# Then run the following commands and adapt the url:
+#
+# .. code-block:: bash
+#
+#     %%bash
+#     pip install gdown  # Install if necessary
+#     gdown "https://drive.google.com/uc?id=1VOV-UM0ws0lopHvoYT_9RAiRdT1y4Kus"
+#     tar -xvf nas_regression.tar.gz
+
+# %% 
+# If you want to remove previously computed results run the following command:
+#
+# .. code-block:: bash
+#
+#     %%bash
+#     rm -rf nas_regression/
+
+# %% 
+# As the search can take some time to finalize we provide a mechanism that checks if results were already computed and skip 
+# the neural architecture search if it is the case.
+max_evals = 250
+
+hpo_results = None
+hpo_results_path = os.path.join(hpo_dir, "results.csv")
+if os.path.exists(hpo_results_path):
+    print("Reloading results...")
+    hpo_results = pd.read_csv(hpo_results_path)
+
+if hpo_results is None or len(hpo_results) < max_evals:
+    print("Running neural architecture search...")
+    hpo_results = run_neural_architecture_search(problem, max_evals)
+
+# %%
+# Analysis of the results
+# -----------------------
+# 
+# We will now look at the results of the search globally in term of evolution of the objective and worker's activity.
+
+from deephyper.analysis.hpo import plot_search_trajectory_single_objective_hpo
+from deephyper.analysis.hpo import plot_worker_utilization
+
+
+fig, axes = plt.subplots(
+    nrows=2,
+    ncols=1,
+    sharex=True,
+    figsize=(WIDTH_PLOTS, HEIGHT_PLOTS),
+)
+
+_ = plot_search_trajectory_single_objective_hpo(
+    hpo_results,
+    mode="min",
+    x_units="seconds",
+    ax=axes[0],
+)
+axes[0].set_yscale("log")
+
+_ = plot_worker_utilization(
+    hpo_results,
+    profile_type="submit/gather",
+    ax=axes[1],
+)
+
+plt.tight_layout()
+
+# %%
+# Then, we split results between successful and failed results if there are some.
+from deephyper.analysis.hpo import filter_failed_objectives
+
+
+hpo_results, hpo_results_failed = filter_failed_objectives(hpo_results)
+
+hpo_results
+
+
+# %%
+# We look at the learning curves of the best model and observe improvements in both training and validation loss:
+
+# .. dropdown: Make learning curves plot
+x_values = np.arange(1, len(baseline_results["metadata"]["train_loss"]) + 1)
+x_min, x_max = x_values.min(), x_values.max()
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+_ = plt.plot(
+    x_values,
+    baseline_results["metadata"]["train_loss"],
+    linestyle=":",
+    label="Baseline Training",
+)
+_ = plt.plot(
+    x_values,
+    baseline_results["metadata"]["val_loss"],
+    linestyle=":",
+    label="Baseline Validation",
+)
+
+i_max = hpo_results["objective"].argmax()
+train_loss = json.loads(hpo_results.iloc[i_max]["m:train_loss"])
+val_loss = json.loads(hpo_results.iloc[i_max]["m:val_loss"])
+x_values = np.arange(1, len(train_loss) + 1)
+x_max = max(x_max, x_values.max())
+_ = plt.plot(
+    x_values,
+    train_loss,
+    alpha=0.8,
+    linestyle="--",
+    label="Best Training",
+)
+_ = plt.plot(
+    x_values,
+    val_loss,
+    alpha=0.8,
+    linestyle="--",
+    label="Best Validation",
+)
+_ = plt.xlim(x_min, x_max)
+_ = plt.grid(which="both", linestyle=":")
+_ = plt.legend()
+_ = plt.xlabel("Epochs")
+_ = plt.ylabel("NLL")
+
+
+# %%
+# Finally, we look at predictions of this best model and observe that it manage to predict much better than the baseline one the right range. 
+from deephyper.analysis.hpo import parameters_from_row
+
+
+hpo_dir = "nas_regression"
+model_checkpoint_dir = os.path.join(hpo_dir, "models")
+job_id = hpo_results.iloc[i_max]["job_id"]
+file_name = f"model_0.{job_id}.pt"
+
+weights_path = os.path.join(model_checkpoint_dir, file_name)
+parameters = parameters_from_row(hpo_results.iloc[i_max])
+
+torch_module = create_model(parameters, y_mu, y_std)
+
+torch_module.load_state_dict(torch.load(weights_path, weights_only=True))
+torch_module.eval()
+
+y_pred = torch_module.forward(to_torch(test_X))
+y_pred_mean = to_numpy(y_pred.loc)
+y_pred_std = to_numpy(y_pred.scale)
+
+# %%
+
+# .. dropdown:: Make prediction plot
+kappa = 1.96
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+_ = plt.scatter(train_X, train_y, s=5, label="Training")
+_ = plt.scatter(valid_X, valid_y, s=5, label="Validation")
+_ = plt.plot(test_X, test_y, linestyle="--", color="gray", label="Test")
+
+_ = plt.plot(test_X, y_pred_mean, label=r"$\mu(x)$")
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred_mean - kappa * y_pred_std).reshape(-1),
+    (y_pred_mean + kappa * y_pred_std).reshape(-1),
+    alpha=0.25,
+    label=r"$\sigma_\text{al}(x)$",
+)
+_ = plt.fill_between([-30, -15], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.fill_between([15, 30], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.xlim(-x_lim, x_lim)
+_ = plt.ylim(-y_lim, y_lim)
+_ = plt.legend(ncols=2)
+_ = plt.xlabel(r"$x$")
+_ = plt.ylabel(r"$f(x)$")
+_ = plt.grid(which="both", linestyle=":")
+
+# %% 
+# Deep ensemble
+# -------------
+#
+# After running the neural architecture search we have an available library of checkpointed models.
+# From this section, you will learn how to combine these models to form an ensemble that can improve both accuracy and provide disentangled uncertainty quantification.
+#
+# We start by importing classes from :mod:`deephyper.predictor` and :mod:`deephyper.ensemble`.
+#
+# The :mod:`deephyper.predictor` module includes subclasses of :class:`deephyper.predictor.Predictor` to wrap predictive models ready for inference. In our case, we will use :class:`deephyper.predictor.torch.TorchPredictor`.
+# The :mod:`deephyper.ensemble` module includes modular components to build an ensemble of predictive models.
+# The ensemble module is organized around loss functions, aggregation functions and selection algorithms.
+# The implementation of these functions is based on Numpy.
+# In this example, we start by wrapping our torch module within a subclass of :class:`deephyper.predictor.torch.TorchPredictor` that we call ``NormalTorchPredictor``. This predictor class is used to make a torch module compatible with our Numpy-based implementation for ensembles.
+#
+# The ``pre_process_inputs`` is used to map a Numpy array to a Torch tensor.
+# The ``post_process_predictions`` is used to map a Torch tensor to a Numpy array.
+# It also formats the prediction as a dictionnary with ``"loc"`` (for the predictive mean) and ``"scale"`` (for the predictive standard deviation) that is necessary for our aggregation function ``MixedNormalAggregator``.
+from deephyper.ensemble import EnsemblePredictor
+from deephyper.ensemble.aggregator import MixedNormalAggregator
+from deephyper.ensemble.loss import NormalNegLogLikelihood
+from deephyper.ensemble.selector import GreedySelector, TopKSelector
+from deephyper.predictor.torch import TorchPredictor
+
+class NormalTorchPredictor(TorchPredictor):
+    def __init__(self, torch_module):
+        super().__init__(torch_module.to(device=device, dtype=dtype))
+
+    def pre_process_inputs(self, X):
+        return to_torch(X)
+
+    def post_process_predictions(self, y):
+        return {
+            "loc": to_numpy(y.loc),
+            "scale": to_numpy(y.scale),
+        }
+
+
+# %%
+# After defining the predictor, we load the checkpointed models to collect their predictions into ``y_predictors``.
+# These predictions are the inputs of our loss, aggregation and selection functions.
+# We also collect the job ids of the checkpointed models into ``job_id_predictors``.
+model_checkpoint_dir = os.path.join(hpo_dir, "models")
+
+y_predictors = []
+job_id_predictors = []
+
+for file_name in tqdm(os.listdir(model_checkpoint_dir)):
+    if not file_name.endswith(".pt"):
+        continue
+
+    weights_path = os.path.join(model_checkpoint_dir, file_name)
+    job_id = int(file_name[6:-3].split(".")[-1])
+
+    row = hpo_results[hpo_results["job_id"] == job_id]
+    if len(row) == 0:
+        continue
+    assert len(row) == 1
+
+    row = row.iloc[0]
+    parameters = parameters_from_row(row)
+    torch_module = create_model(parameters, y_mu, y_std)
+    try:
+        torch_module.load_state_dict(torch.load(weights_path, weights_only=True))
+    except RuntimeError:
+        continue
+
+    predictor = NormalTorchPredictor(torch_module)
+    y_pred = predictor.predict(valid_X)
+    y_predictors.append(y_pred)
+    job_id_predictors.append(job_id)
+
+# %%
+# Ensemble selection
+# ------------------
+#
+# This is where the ensemble selection logic happens. We use the :class:`deephyper.ensemble.selector.GreedySelector` or :class:`deephyper.ensemble.selector.TopKSelector` class.
+# The top-k selection, selects the topk-k models according to the given ``los_func`` and weight them equally in the ensemble.
+# The greedy selection, iteratively selects models from the checkpoints that improves the current ensemble.
+#
+# The ``aggregator`` is the logic that combines a set of predictors into a single predictor to form the ensemble's prediction.
+# In our case, we use the :class:`deephyper.ensemble.aggregator.MixedNormalAggregator` that approximates a mixture of normal distribution (each normal distribution is the output of a checkpointed model) as a normal distribution.
+#
+# To try top-k or greedy selection just uncomment/comment the corresponding code.
+# This part of the code is fast to compute.
+k = 50
+
+# Top-K Selection
+# selector = TopKSelector(
+#     loss_func=NormalNegLogLikelihood(),
+#     k=k,
+# )
+
+# Greedy Selection
+selector = GreedySelector(
+    loss_func=NormalNegLogLikelihood(),
+    aggregator=MixedNormalAggregator(),
+    k=k,
+    max_it=k,
+    k_init=3,
+    early_stopping=True,
+    with_replacement=True,
+    bagging=True,
+    verbose=True,
+)
+
+selected_predictors_indexes, selected_predictors_weights = selector.select(
+    valid_y,
+    y_predictors,
+)
+
+print(f"{selected_predictors_indexes=}")
+print(f"{selected_predictors_weights=}")
+
+selected_predictors_job_ids = np.array(job_id_predictors)[selected_predictors_indexes]
+selected_predictors_job_ids
+
+print(f"{selected_predictors_job_ids=}")
+
+# %%
+# Evaluation of the ensemble
+# --------------------------
+#
+# Now that we have a set of predictors with their corresponding weights in the ensemble we can look at the predictions.
+# For this, we use the :class:`deephyper.ensemble.EnsemblePredictor` class.
+# This class can use the :class:`deephyper.evaluator.Evaluator` to parallelize the inference of ensemble members.
+# Then, we need to give it the list of ``predictors``, ``weights`` and the ``aggregator``.
+# For inference, we set ``decomposed_scale=True`` for the :class:`deephyper.ensemble.aggregator.MixedNormalAggregator` as we want
+# to predict disentangled epistemic and aleatoric uncertainty using the law of total variance:
+#
+# .. math::
+# 
+#     V_Y[Y|X=x] = \underbrace{E_\Theta\left[V_Y[Y|X=x;\Theta\right]}_\text{Aleatoric Uncertainty} + \underbrace{V_\Theta\left[E_Y[Y|X=x;\Theta]\right]}_\text{Epistemic Uncertainty}
+# 
+# where :math:`\Theta` is the random variable that represents a concatenation of weights and hyperparameters, :math:`Y`` is the random variable representing a target prediction, and :math:`X` is the random variable representing an observed input.
+predictors = []
+
+hpo_dir = "nas_regression"
+model_checkpoint_dir = os.path.join(hpo_dir, "models")
+
+for job_id in selected_predictors_job_ids:
+    file_name = f"model_0.{job_id}.pt"
+
+    weights_path = os.path.join(model_checkpoint_dir, file_name)
+
+    row = hpo_results[hpo_results["job_id"] == job_id].iloc[0]
+    parameters = parameters_from_row(row)
+    torch_module = create_model(parameters, y_mu, y_std)
+    torch_module.load_state_dict(torch.load(weights_path, weights_only=True))
+    predictor = NormalTorchPredictor(torch_module)
+    predictors.append(predictor)
+
+ensemble = EnsemblePredictor(
+    predictors=predictors,
+    weights=selected_predictors_weights,
+    aggregator=MixedNormalAggregator(decomposed_scale=True),
+)
+
+y_pred = ensemble.predict(test_X)
+
+# %%
+#
+# In the visualization, we can first observe that the mean prediction is close to the true function.
+#
+# Then, to visualize both uncertainties together we plot the variance.
+# The goal is to observe the epistemic component vanish in areas where we observed data.
+
+# .. dropdown:: Make uncertainty plot
+# sphinx_gallery_thumbnail_number = 7
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+_ = plt.scatter(train_X, train_y, s=5, label="Training")
+_ = plt.scatter(valid_X, valid_y, s=5, label="Validation")
+_ = plt.plot(test_X, test_y, linestyle="--", color="gray", label="Test")
+_ = plt.plot(test_X, y_pred["loc"], label=r"$\mu(x)$")
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred["loc"] - y_pred["scale_aleatoric"]**2).reshape(-1),
+    (y_pred["loc"] + y_pred["scale_aleatoric"]**2).reshape(-1),
+    alpha=0.25,
+    label=r"$\sigma_\text{al}^2(x)$",
+)
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred["loc"] - y_pred["scale_aleatoric"]**2).reshape(-1),
+    (y_pred["loc"] - y_pred["scale_aleatoric"]**2 - y_pred["scale_epistemic"]**2).reshape(-1),
+    alpha=0.25,
+    color="red",
+    label=r"$\sigma_\text{ep}^2(x)$",
+)
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred["loc"] + y_pred["scale_aleatoric"]**2).reshape(-1),
+    (y_pred["loc"] + y_pred["scale_aleatoric"]**2 + y_pred["scale_epistemic"]**2).reshape(-1),
+    alpha=0.25,
+    color="red",
+)
+_ = plt.fill_between([-30, -15], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.fill_between([15, 30], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.xlim(-x_lim, x_lim)
+_ = plt.ylim(-y_lim, y_lim)
+_ = plt.legend(ncols=2)
+_ = plt.xlabel(r"$x$")
+_ = plt.ylabel(r"$f(x)$")
+_ = plt.grid(which="both", linestyle=":")
+
+# %%
+# Aleatoric Uncertainty
+# ~~~~~~~~~~~~~~~~~~~~~
+#
+# Now, if we isolate the aleatoric uncertainty we observe that we somewhat correctly estimated the lower aleatoric uncertainty on the left side, and larger on the right side.
+
+# .. dropdown:: Make aleatoric uncertainty plot
+kappa = 1.96
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+_ = plt.scatter(train_X, train_y, s=5, label="Training")
+_ = plt.scatter(valid_X, valid_y, s=5, label="Validation")
+_ = plt.plot(test_X, test_y, linestyle="--", color="gray", label="Test")
+_ = plt.plot(test_X, y_pred["loc"], label=r"$\mu(x)$")
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred["loc"] - kappa * y_pred["scale_aleatoric"]).reshape(-1),
+    (y_pred["loc"] + kappa * y_pred["scale_aleatoric"]).reshape(-1),
+    alpha=0.25,
+    label=r"$\sigma_\text{al}(x)$",
+)
+_ = plt.fill_between([-30, -15], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.fill_between([15, 30], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.xlim(-x_lim, x_lim)
+_ = plt.ylim(-y_lim, y_lim)
+_ = plt.legend(ncols=2)
+_ = plt.xlabel(r"$x$")
+_ = plt.ylabel(r"$f(x)$")
+_ = plt.grid(which="both", linestyle=":")
+
+# %%
+# Epistemic uncertainty
+# ~~~~~~~~~~~~~~~~~~~~~
+#
+# Finally, if we isole the epistemic uncertainty we observe that it vanishes in the grey areas where we observed data and grows in areas were we did not have data.
+
+# .. dropdown:: Make epistemic uncertainty plot
+kappa = 1.96
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+_ = plt.scatter(train_X, train_y, s=5, label="Training")
+_ = plt.scatter(valid_X, valid_y, s=5, label="Validation")
+_ = plt.plot(test_X, test_y, linestyle="--", color="gray", label="Test")
+_ = plt.plot(test_X, y_pred["loc"], label=r"$\mu(x)$")
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred["loc"] - kappa * y_pred["scale_epistemic"]).reshape(-1),
+    (y_pred["loc"] + kappa * y_pred["scale_epistemic"]).reshape(-1),
+    alpha=0.25,
+    color="red",
+    label=r"$\sigma_\text{ep}(x)$",
+)
+_ = plt.fill_between([-30, -15], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.fill_between([15, 30], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.xlim(-x_lim, x_lim)
+_ = plt.ylim(-y_lim, y_lim)
+_ = plt.legend(ncols=2)
+_ = plt.xlabel(r"$x$")
+_ = plt.ylabel(r"$f(x)$")
+_ = plt.grid(which="both", linestyle=":")
+
+# %%
+# After defining the predictor, we load the checkpointed models to collect their predictions into ``y_predictors``.
+# These predictions are the inputs of our loss, aggregation and selection functions.
+# We also collect the job ids of the checkpointed models into ``job_id_predictors``.
+model_checkpoint_dir = os.path.join(hpo_dir, "models")
+
+predictors = []
+
+for file_name in tqdm(os.listdir(model_checkpoint_dir)):
+    if not file_name.endswith(".pt"):
+        continue
+
+    weights_path = os.path.join(model_checkpoint_dir, file_name)
+    job_id = int(file_name[6:-3].split(".")[-1])
+
+    row = hpo_results[hpo_results["job_id"] == job_id]
+    if len(row) == 0:
+        continue
+    assert len(row) == 1
+
+    row = row.iloc[0]
+    parameters = parameters_from_row(row)
+    torch_module = create_model(parameters, y_mu, y_std)
+    try:
+        torch_module.load_state_dict(torch.load(weights_path, weights_only=True))
+    except RuntimeError:
+        continue
+
+    predictor = NormalTorchPredictor(torch_module)
+    predictors.append(predictor)
+
+# %% 
+# Error Model Greedy Selection
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# Greedy Selection
+
+from deephyper.ensemble._erm_ensemble import ERMEnsemblePredictor
+
+ensemble = ERMEnsemblePredictor(
+    predictors=predictors,
+    aggregator=MixedNormalAggregator(decomposed_scale=True),
+)
+
+ensemble.fit(valid_X, valid_y)
+
+y_pred = ensemble.predict(test_X)
+
+# %%
+#
+# In the visualization, we can first observe that the mean prediction is close to the true function.
+#
+# Then, to visualize both uncertainties together we plot the variance.
+# The goal is to observe the epistemic component vanish in areas where we observed data.
+
+# .. dropdown:: Make uncertainty plot
+# sphinx_gallery_thumbnail_number = 7
+plt.close("all")
+_ = plt.figure(figsize=(WIDTH_PLOTS, HEIGHT_PLOTS))
+_ = plt.scatter(train_X, train_y, s=5, label="Training")
+_ = plt.scatter(valid_X, valid_y, s=5, label="Validation")
+_ = plt.plot(test_X, test_y, linestyle="--", color="gray", label="Test")
+_ = plt.plot(test_X, y_pred["loc"], label=r"$\mu(x)$")
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred["loc"] - y_pred["scale_aleatoric"]**2).reshape(-1),
+    (y_pred["loc"] + y_pred["scale_aleatoric"]**2).reshape(-1),
+    alpha=0.25,
+    label=r"$\sigma_\text{al}^2(x)$",
+)
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred["loc"] - y_pred["scale_aleatoric"]**2).reshape(-1),
+    (y_pred["loc"] - y_pred["scale_aleatoric"]**2 - y_pred["scale_epistemic"]**2).reshape(-1),
+    alpha=0.25,
+    color="red",
+    label=r"$\sigma_\text{ep}^2(x)$",
+)
+_ = plt.fill_between(
+    test_X.reshape(-1),
+    (y_pred["loc"] + y_pred["scale_aleatoric"]**2).reshape(-1),
+    (y_pred["loc"] + y_pred["scale_aleatoric"]**2 + y_pred["scale_epistemic"]**2).reshape(-1),
+    alpha=0.25,
+    color="red",
+)
+_ = plt.fill_between([-30, -15], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.fill_between([15, 30], [-y_lim, -y_lim], [y_lim, y_lim], color="gray", alpha=0.25)
+_ = plt.xlim(-x_lim, x_lim)
+_ = plt.ylim(-y_lim, y_lim)
+_ = plt.legend(ncols=2)
+_ = plt.xlabel(r"$x$")
+_ = plt.ylabel(r"$f(x)$")
+_ = plt.grid(which="both", linestyle=":")
+plt.show()
\ No newline at end of file
diff --git a/src/deephyper/ensemble/_erm_ensemble.py b/src/deephyper/ensemble/_erm_ensemble.py
new file mode 100644
index 00000000..4ebda1b7
--- /dev/null
+++ b/src/deephyper/ensemble/_erm_ensemble.py
@@ -0,0 +1,204 @@
+import copy
+from typing import Dict, Sequence
+
+import numpy as np
+from tqdm import tqdm
+
+from deephyper.ensemble.aggregator import Aggregator
+from deephyper.ensemble.loss import NormalNegLogLikelihood 
+from deephyper.ensemble.selector._erm_greedy import ERMGreedySelector
+from deephyper.evaluator import Evaluator, RunningJob
+from deephyper.evaluator.callback import TqdmCallback
+from deephyper.evaluator.storage import NullStorage
+from deephyper.predictor import Predictor, PredictorLoader
+
+
+def predict_with_predictor(predictor: Predictor | PredictorLoader, X: np.ndarray):
+    if isinstance(predictor, PredictorLoader):
+        predictor = predictor.load()
+    return predictor.predict(X)
+
+
+def _wrapper_predict_with_predictor(job: RunningJob):
+    try:
+        return predict_with_predictor(**job.parameters)
+    except Exception as exception:
+        return exception
+
+
+class ERMEnsemblePredictor(Predictor):
+    """A predictor that is itself an ensemble of multiple predictors.
+
+    Args:
+        predictors (Sequence[Predictor | PredictorLoader]): the list of predictors to put in the
+            ensemble. The sequence can be composed of ``Predictor`` (i.e., the model is already
+            loaded in memory) or ``PredictorLoader`` to perform the loading remotely and in
+            parallel. In the later case, the ``.load()`` function is called for each inference.
+
+        aggregator (Aggregator): the aggregation function to fuse the predictions of the predictors
+            into one prediction.
+
+        weights (Sequence[float], optional): the weights of the predictors in the aggregation.
+            Defaults to ``None``.
+
+        evaluator (str | Dict, optional): The parallel strategy to compute predictions from the
+            list of predictions. If it is a ``str`` it must be a possible ``method`` of
+            ``Evaluator.create(..., method=...)``. If it is a ``dict`` it must have two keys
+            ``method`` and ``method_kwargs`` such as ``Evaluator.create(...)``. Defaults to
+            ``None`` which is equivalent to ``evaluator="serial"`` for serial evaluations.
+
+    Raises:
+        ValueError: when the type of the ``evaluator`` argument is not ``str`` or ``dict``.
+    """
+
+    def __init__(
+        self,
+        predictors: Sequence[Predictor | PredictorLoader],
+        aggregator: Aggregator,
+        weights: Sequence[float] = None,
+        evaluator: str | Dict = None,
+    ):
+        self.predictors = predictors
+        self.aggregator = aggregator
+        self.selector = ERMGreedySelector(
+            loss_func=NormalNegLogLikelihood(),
+            aggregator=aggregator,
+            k=10,
+            max_it=10,
+            k_init=3,
+            early_stopping=False,
+            with_replacement=True,
+            bagging=True,
+            verbose=1,
+        )
+        self.weights = weights
+        self._evaluator = None
+
+        if evaluator is None:
+            self.evaluator_method = "thread"
+            self.evaluator_method_kwargs = {}
+        elif isinstance(evaluator, str):
+            self.evaluator_method = evaluator
+            self.evaluator_method_kwargs = {}
+        elif isinstance(evaluator, dict):
+            self.evaluator_method = evaluator.get("method", "serial")
+            self.evaluator_method_kwargs = evaluator.get("method_kwargs", {})
+        else:
+            raise ValueError(f"evaluator must be either None or str or dict, got {type(evaluator)}")
+        self.init_evaluator()
+
+    def init_evaluator(self):
+        """Initialize an evaluator for the ensemble.
+
+        Returns:
+            Evaluator: An evaluator instance.
+        """
+        method_kwargs = {
+            "storage": NullStorage(),
+            "run_function_kwargs": {},
+        }
+        method_kwargs.update(self.evaluator_method_kwargs)
+        self._evaluator = Evaluator.create(
+            run_function=_wrapper_predict_with_predictor,
+            method=self.evaluator_method,
+            method_kwargs=method_kwargs,
+        )
+
+    def fit(self, X: np.ndarray, y: np.ndarray):
+        y_predictors = self.predictions_from_predictors(X, self.predictors)
+        self.selector.fit(X, y, y_predictors)
+
+    def predict(self, X: np.ndarray):
+        """Compute the prediction of the ensemble.
+
+        Args:
+            X (np.ndarray): the input query for the prediction.
+
+        Returns:
+            np.ndarray: the target prediction.
+        """
+        y_predictors = self.predictions_from_predictors(X, self.predictors)
+        y_dummy = self.aggregator.aggregate(y_predictors[:1])
+        if isinstance(y_dummy, dict):
+            for k in y_dummy:
+                y_dummy[k] = y_dummy[k][:1]
+            y = {k: [] for k in y_dummy}
+        else:
+            y_dummy = y_dummy[:1]
+            y = []
+
+        for i in tqdm(range(len(X))):
+            X_i = X[i : i + 1]
+            if isinstance(y_predictors[0], dict):
+                y_predictors_i = [{k: y_[k][i : i + 1] for k in y_} for y_ in y_predictors]
+            else:
+                y_predictors_i = [y_[i : i + 1] for y_ in y_predictors]
+            indexes, weights = self.selector.select(X_i, y_predictors_i)
+
+            y_predictors_i = [y_predictors_i[j] for j in indexes]
+
+            y_i = self.aggregator.aggregate(y_predictors_i, weights=weights)
+            if isinstance(y, dict):
+                for k in y:
+                    y[k].append(y_i[k][0])
+            else:
+                y.append(y_i[0])
+
+        if isinstance(y, dict):
+            for k in y:
+                y[k] = np.asarray(y[k])
+        else:
+            y = np.asarray(y[k])
+
+        return y
+
+    def predictions_from_predictors(
+        self, X: np.ndarray, predictors: Sequence[Predictor | PredictorLoader]
+    ):
+        """Compute the predictions of a list of predictors.
+
+        Args:
+            X (np.ndarray): the input query for the predictions.
+            predictors (Sequence[Predictor]): the list of predictors to compute the predictions.
+
+        Returns:
+            List[np.ndarray]: the sequence of predictions in the same order that the list of
+            predictors.
+        """
+        n_jobs_submitted = len(predictors)
+
+        for cb in self._evaluator._callbacks:
+            if isinstance(cb, TqdmCallback):
+                cb.set_max_evals(n_jobs_submitted)
+
+        self._evaluator.submit(
+            [
+                {
+                    "predictor": predictor,
+                    "X": X,
+                }
+                for predictor in predictors
+            ]
+        )
+
+        jobs_done = []
+        while len(jobs_done) != n_jobs_submitted:
+            new_jobs_done = self._evaluator.gather("BATCH", size=1)
+            jobs_done.extend(new_jobs_done)
+        jobs_done = list(sorted(jobs_done, key=lambda j: int(j.id.split(".")[-1])))
+
+        self._evaluator.close()
+
+        y_pred = []
+        for i, job in enumerate(jobs_done):
+            if isinstance(job.output, Exception):
+                try:
+                    raise job.output
+                except Exception:
+                    raise RuntimeError(
+                        f"Failed to call .predict(X) with predictors[{i}]: {predictors[i]}"
+                    )
+            else:
+                y_pred.append(job.output)
+
+        return y_pred
diff --git a/src/deephyper/ensemble/selector/_erm_greedy.py b/src/deephyper/ensemble/selector/_erm_greedy.py
new file mode 100644
index 00000000..d8b82bc3
--- /dev/null
+++ b/src/deephyper/ensemble/selector/_erm_greedy.py
@@ -0,0 +1,237 @@
+from typing import Callable, List
+
+import numpy as np
+from sklearn.ensemble import RandomForestRegressor, ExtraTreesRegressor
+from sklearn.linear_model import LinearRegression, Ridge
+from sklearn.metrics import r2_score
+from sklearn.model_selection import train_test_split
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import PolynomialFeatures
+from sklearn.preprocessing import QuantileTransformer
+from sklearn.utils import check_random_state
+from sklearn.compose import TransformedTargetRegressor
+from sklearn.neural_network import MLPRegressor
+
+from deephyper.ensemble.aggregator._aggregator import Aggregator
+from deephyper.ensemble.selector._selector import Selector
+
+
+class ERMGreedySelector(Selector):
+    """Selection method implementing Greedy (a.k.a., Caruana) selection.
+
+    This method iteratively and greedily selects the predictors that minimize
+    the loss when aggregated together.
+
+    Args:
+        loss_func (Callable or Loss): a loss function that takes two arguments: the true target
+            values and the predicted target values.
+        aggregator (Aggregator): The aggregator to use to combine the predictions of the selected
+            predictors.
+        k (int, optional): The number of unique predictors to select for the ensemble. Defaults to
+            ``5``.
+        k_init (int, optional): Regularization parameter for greedy selection. It is the number of
+            predictors to select in the initialization step. Defaults to ``1``.
+        max_it (int, optional): Maximum number of iterations which also corresponds to the number
+            of non-unique predictors added to the ensemble. Defaults to ``-1``.
+        eps_tol (float, optional): Tolerance for the stopping criterion. Defaults to ``1e-3``.
+        with_replacement (bool, optional): Performs greedy selection with replacement of models
+            already selected. Defaults to ``True``.
+        early_stopping (bool, optional): Stops the ensemble selection as soon as the loss stops
+            improving. Defaults to ``True``.
+        bagging (bool, optional): Performs boostrap resampling of available predictors at each
+            iteration. This can be particularly useful when the dataset used for selection is
+            small. Defaults to ``False``.
+        verbose (bool, optional):
+            Turns on the verbose mode. Defaults to ``False``.
+    """
+
+    def __init__(
+        self,
+        loss_func: Callable,
+        aggregator: Aggregator,
+        k: int = 5,
+        k_init: int = 5,
+        max_it: int = -1,
+        eps_tol: float = 1e-3,
+        with_replacement: bool = True,
+        early_stopping: bool = True,
+        bagging: bool = False,
+        random_state=None,
+        verbose: bool = False,
+    ):
+        super().__init__(loss_func)
+        self.aggregator = aggregator
+        self.k = k
+        self.k_init = k_init
+        self.max_it = max_it
+        self.eps_tol = eps_tol
+        self.with_replacement = with_replacement
+        self.early_stopping = early_stopping
+        self.bagging = bagging
+        self.random_state = check_random_state(random_state)
+        self.verbose = verbose
+
+        # self.erm_model = ExtraTreesRegressor(min_samples_leaf=4)
+        # self.erm_model = LinearRegression()
+        # self.erm_model = Pipeline(
+        #     [
+        #         ("poly", PolynomialFeatures(degree=2)),
+        #         ("linear", Ridge(fit_intercept=False)),
+        #     ]
+        # )
+        self.erm_model = MLPRegressor(
+            hidden_layer_sizes=(100, 100,),
+            learning_rate_init=0.001,
+            max_iter=10_000,
+            verbose=True,
+            n_iter_no_change=100,
+        )
+        self.erm_model = TransformedTargetRegressor(
+            regressor=self.erm_model, transformer=QuantileTransformer(output_distribution="uniform")
+        )
+
+    def _aggregate(self, y_predictors: np.ndarray, weights: List = None):
+        return self.aggregator.aggregate(y_predictors, weights)
+
+    def fit(self, X, y, y_predictors):
+        # Creates the Error model
+        # Inputs: (X, y)
+        # Output: Loss(y)
+        if isinstance(y_predictors[0], dict):
+            X_erm = np.concatenate(
+                [
+                    np.concatenate((X, y_pred_i["loc"], y_pred_i["scale"]), axis=1)
+                    for y_pred_i in y_predictors
+                ],
+                axis=0,
+            )
+        else:
+            X_erm = np.concatenate(
+                [np.concatenate((X, y_pred_i), axis=1) for y_pred_i in y_predictors],
+                axis=0,
+            )
+        y_erm = np.concatenate(
+            [self.loss_func(y, y_pred_i) for y_pred_i in y_predictors],
+            axis=0,
+        ).reshape(-1)
+        print("min:", np.min(y_erm))
+        print("max:", np.max(y_erm))
+        # print(f"{np.shape(X_erm)=}")
+        # print(f"{np.shape(y_erm)=}")
+        X_train, X_valid, y_train, y_valid = train_test_split(X_erm, y_erm, test_size=0.33)
+        q1 = np.quantile(y_erm, q=0.25)
+        q2 = np.quantile(y_erm, q=0.5)
+        q3 = np.quantile(y_erm, q=0.75)
+        print(f"{q1:.2f} | {q2:.2f} | {q3:.3f}")
+        threshold = q3 + 1.5 * (q3 - q1)
+        y_train[y_train > threshold] = threshold
+        y_valid[y_valid > threshold] = threshold
+        self.erm_model.fit(X_train, y_train)
+        r2_score_train = self.erm_model.score(X_train, y_train)
+        r2_score_valid = self.erm_model.score(X_valid, y_valid)
+        print(f"{r2_score_train=:.3f}")
+        print(f"{r2_score_valid=:.3f}")
+
+        # import matplotlib.pyplot as plts
+
+        # plt.close("all")
+        # plt.figure()
+        # plt.boxplot(y_erm)
+        # plt.show()
+
+        # y_valid_pred = self.erm_model.predict(X_valid)
+        # plt.figure()
+        # plt.scatter(y_valid, y_valid_pred)
+        # plt.xlabel("True")
+        # plt.ylabel("Pred")
+        # plt.show()
+
+    def _evaluate(self, X, y_pred) -> float:
+        if isinstance(y_pred, dict):
+            if "scale_aleatoric" in y_pred and "scale_epistemic" in y_pred:
+                scale = np.sqrt(y_pred["scale_aleatoric"] ** 2 + y_pred["scale_epistemic"] ** 2)
+                X_erm = np.concatenate((X, y_pred["loc"], scale), axis=1)
+            else:
+                X_erm = np.concatenate((X, y_pred["loc"], y_pred["scale"]), axis=1)
+        else:
+            X_erm = np.concatenate((X, y_pred), axis=1)
+        return self._reduce(self.erm_model.predict(X_erm))
+
+    def select(self, X, y_predictors) -> tuple[list[int], list[float]]:
+        # Initialization
+        losses = [self._evaluate(X, y_pred_i) for y_pred_i in y_predictors]
+        selected_indices = np.argsort(losses)[: self.k_init].tolist()
+        selected_indices_weights = [1 / self.k_init] * self.k_init
+        loss_min = self._evaluate(X, self._aggregate([y_predictors[i] for i in selected_indices]))
+        n_predictors = len(y_predictors)
+        bagged_predictors = None
+
+        if self.verbose:
+            tmp = [losses[i] for i in selected_indices]
+            print(f"Ensemble initialized with {selected_indices} with loss {tmp}")
+
+        # Greedy steps
+        it = 0
+        while (self.max_it < 0 or it < self.max_it) and len(np.unique(selected_indices)) < self.k:
+            losses = []
+
+            if self.bagging:
+                bagged_predictors = np.unique(
+                    self.random_state.randint(low=0, high=n_predictors, size=n_predictors)
+                )
+
+            for i in range(n_predictors):
+                # Applying conditions that ignore some indices in the selection
+                if len(selected_indices) == 1 and i in selected_indices:
+                    losses.append(np.nan)
+                    continue
+
+                if not self.with_replacement and i in selected_indices:
+                    losses.append(np.nan)
+                    continue
+
+                if self.bagging and i not in bagged_predictors:
+                    losses.append(np.nan)
+                    continue
+
+                indices_ = selected_indices + [i]
+                indices_, indices_weights_ = np.unique(indices_, return_counts=True)
+                indices_weights_ = indices_weights_ / np.sum(indices_weights_)
+                y_ = [y_predictors[i] for i in indices_]
+                score = self._evaluate(
+                    X,
+                    self._aggregate(y_, indices_weights_),
+                )
+                losses.append(score)
+
+            i_min_ = int(np.nanargmin(losses))
+            loss_min_ = losses[i_min_]
+            it += 1
+
+            # The second condition is related to numerical errors
+            if (self.early_stopping and loss_min_ >= (loss_min - self.eps_tol)) or (
+                len(np.unique(selected_indices)) == 1 and selected_indices[0] == i_min_
+            ):
+                if self.verbose:
+                    print(f"Step {it}, ensemble selection stopped")
+                break
+
+            loss_min = loss_min_
+            selected_indices.append(i_min_)
+
+            if self.verbose:
+                print(
+                    f"Step {it}, ensemble is {selected_indices}, new member {i_min_} with"
+                    f" loss {loss_min}"
+                )
+
+        selected_indices, selected_indices_weights = np.unique(selected_indices, return_counts=True)
+        selected_indices_weights = selected_indices_weights / np.sum(selected_indices_weights)
+
+        if self.verbose:
+            print(
+                f"After {it} steps, the final ensemble is {selected_indices} with "
+                f"weights {selected_indices_weights}"
+            )
+
+        return selected_indices.tolist(), selected_indices_weights.tolist()