Repository: vanderschaarlab/temporai (language: python)

Issue:
Make metrics plugin-able



Discussion / hints:

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index b5940265..a51cfaba 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -71,7 +71,7 @@ repos:
   - id: flake8
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.6.1
+  rev: v1.7.0
   hooks:
   - id: mypy
     files: "^src/"
diff --git a/README.md b/README.md
index e8816c6c..ca686297 100644
--- a/README.md
+++ b/README.md
@@ -472,6 +472,7 @@ Prediction where targets are temporal (time series).
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial01_custom_method.ipynb) - [Writing a Custom Method Plugin](./tutorials/extending/tutorial01_custom_method.ipynb)
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial02_testing_custom_method.ipynb) - [Testing a Custom Method Plugin](./tutorials/extending/tutorial02_testing_custom_method.ipynb)
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial03_custom_datasource.ipynb) - [Writing a Custom Data Source Plugin](./tutorials/extending/tutorial03_custom_datasource.ipynb)
+- [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial04_custom_metric.ipynb) - [Writing a Custom Metric Plugin](./tutorials/extending/tutorial04_custom_metric.ipynb)
 
 
 
diff --git a/docs/overview.md b/docs/overview.md
index dee939a2..992b72c8 100644
--- a/docs/overview.md
+++ b/docs/overview.md
@@ -670,6 +670,7 @@ Prediction where targets are temporal (time series).
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial01_custom_method.ipynb) - [Writing a Custom Method Plugin](https://github.com/vanderschaarlab/temporai/tree/main/tutorials/extending/tutorial01_custom_method.ipynb)
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial02_testing_custom_method.ipynb) - [Testing a Custom Method Plugin](https://github.com/vanderschaarlab/temporai/tree/main/tutorials/extending/tutorial02_testing_custom_method.ipynb)
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial03_custom_datasource.ipynb) - [Writing a Custom Data Source Plugin](https://github.com/vanderschaarlab/temporai/tree/main/tutorials/extending/tutorial03_custom_datasource.ipynb)
+- [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial04_custom_metric.ipynb) - [Writing a Custom Metric Plugin](https://github.com/vanderschaarlab/temporai/tree/main/tutorials/extending/tutorial04_custom_metric.ipynb)
 
 
 
diff --git a/docs/user_guide/data/tutorial03_datasources.ipynb b/docs/user_guide/data/tutorial03_datasources.ipynb
index 4e66a146..58e66ff3 100644
--- a/docs/user_guide/data/tutorial03_datasources.ipynb
+++ b/docs/user_guide/data/tutorial03_datasources.ipynb
@@ -64,7 +64,7 @@
     {
      "data": {
       "text/plain": [
-       "<tempor.datasources.prediction.one_off.plugin_sine.SineDataSource at 0x7f18cd1f8670>"
+       "<tempor.datasources.prediction.one_off.plugin_sine.SineDataSource at 0x7ff8f84f6070>"
       ]
      },
      "execution_count": null,
@@ -564,7 +564,7 @@
     {
      "data": {
       "text/plain": [
-       "<tempor.datasources.prediction.one_off.plugin_sine.SineDataSource at 0x7f18cd228cd0>"
+       "<tempor.datasources.prediction.one_off.plugin_sine.SineDataSource at 0x7ff8f8552d90>"
       ]
      },
      "execution_count": null,
@@ -646,7 +646,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "SineDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "OneOffPredictionDataset(\n",
       "    time_series=TimeSeriesSamples([100, *, 5]),\n",
       "    static=StaticSamples([100, 4]),\n",
@@ -1075,7 +1081,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "GoogleStocksDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "OneOffPredictionDataset(\n",
       "    time_series=TimeSeriesSamples([50, *, 5]),\n",
       "    predictive=OneOffPredictionTaskData(targets=StaticSamples([50, 1]))\n",
@@ -1589,24 +1601,24 @@
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th></th>\n",
-       "      <th>post-supper_blood_glucose_measurement</th>\n",
-       "      <th>less-than-usual_meal_ingestion</th>\n",
-       "      <th>regular_insulin_dose</th>\n",
-       "      <th>more-than-usual_exercise_activity</th>\n",
+       "      <th>post-lunch_blood_glucose_measurement</th>\n",
+       "      <th>more-than-usual_meal_ingestion</th>\n",
        "      <th>unspecified_special_event</th>\n",
        "      <th>typical_exercise_activity</th>\n",
-       "      <th>nph_insulin_dose</th>\n",
        "      <th>less-than-usual_exercise_activity</th>\n",
-       "      <th>typical_meal_ingestion</th>\n",
-       "      <th>pre-breakfast_blood_glucose_measurement</th>\n",
-       "      <th>pre-lunch_blood_glucose_measurement</th>\n",
-       "      <th>pre-supper_blood_glucose_measurement</th>\n",
+       "      <th>post-supper_blood_glucose_measurement</th>\n",
+       "      <th>more-than-usual_exercise_activity</th>\n",
        "      <th>unspecified_blood_glucose_measurement</th>\n",
-       "      <th>post-lunch_blood_glucose_measurement</th>\n",
-       "      <th>more-than-usual_meal_ingestion</th>\n",
-       "      <th>post-breakfast_blood_glucose_measurement</th>\n",
-       "      <th>ultralente_insulin_dose</th>\n",
        "      <th>pre-snack_blood_glucose_measurement</th>\n",
+       "      <th>ultralente_insulin_dose</th>\n",
+       "      <th>less-than-usual_meal_ingestion</th>\n",
+       "      <th>pre-supper_blood_glucose_measurement</th>\n",
+       "      <th>regular_insulin_dose</th>\n",
+       "      <th>pre-breakfast_blood_glucose_measurement</th>\n",
+       "      <th>nph_insulin_dose</th>\n",
+       "      <th>post-breakfast_blood_glucose_measurement</th>\n",
+       "      <th>pre-lunch_blood_glucose_measurement</th>\n",
+       "      <th>typical_meal_ingestion</th>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>sample_idx</th>\n",
@@ -1637,19 +1649,19 @@
        "      <th>0</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>9.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>13.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>100.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>9.0</td>\n",
+       "      <td>100.0</td>\n",
+       "      <td>13.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1658,7 +1670,6 @@
        "      <th>1</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1667,8 +1678,9 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>119.0</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>119.0</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1684,12 +1696,12 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>123.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>123.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1700,19 +1712,19 @@
        "      <th>3</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>10.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>13.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>216.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>10.0</td>\n",
+       "      <td>216.0</td>\n",
+       "      <td>13.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1721,7 +1733,6 @@
        "      <th>4</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>2.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1732,6 +1743,7 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>2.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1770,12 +1782,12 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>145.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>145.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1786,19 +1798,19 @@
        "      <th>147</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>1.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>259.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>259.0</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1811,7 +1823,6 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1820,6 +1831,7 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1832,7 +1844,6 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1841,6 +1852,7 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1853,7 +1865,6 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1862,6 +1873,7 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1873,19 +1885,229 @@
       ],
       "text/plain": [
        "TimeSeriesSamples with data:\n",
-       "                     post-supper_blood_glucose_measurement  ...  pre-snack_blood_glucose_measurement\n",
-       "sample_idx time_idx                                         ...                                     \n",
-       "0          0                                           NaN  ...                                  NaN\n",
-       "           1                                           NaN  ...                                  NaN\n",
-       "           2                                           NaN  ...                                  NaN\n",
-       "           3                                           NaN  ...                                  NaN\n",
-       "           4                                           NaN  ...                                  NaN\n",
-       "...                                                    ...  ...                                  ...\n",
-       "69         146                                         NaN  ...                                  NaN\n",
-       "           147                                         NaN  ...                                  NaN\n",
-       "           148                                         NaN  ...                                  NaN\n",
-       "           149                                         NaN  ...                                  NaN\n",
-       "           150                                         NaN  ...                                  NaN\n",
+       "                     post-lunch_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                         \n",
+       "0          0                                          NaN   \n",
+       "           1                                          NaN   \n",
+       "           2                                          NaN   \n",
+       "           3                                          NaN   \n",
+       "           4                                          NaN   \n",
+       "...                                                   ...   \n",
+       "69         146                                        NaN   \n",
+       "           147                                        NaN   \n",
+       "           148                                        NaN   \n",
+       "           149                                        NaN   \n",
+       "           150                                        NaN   \n",
+       "\n",
+       "                     more-than-usual_meal_ingestion  \\\n",
+       "sample_idx time_idx                                   \n",
+       "0          0                                    NaN   \n",
+       "           1                                    NaN   \n",
+       "           2                                    NaN   \n",
+       "           3                                    NaN   \n",
+       "           4                                    NaN   \n",
+       "...                                             ...   \n",
+       "69         146                                  NaN   \n",
+       "           147                                  NaN   \n",
+       "           148                                  NaN   \n",
+       "           149                                  NaN   \n",
+       "           150                                  NaN   \n",
+       "\n",
+       "                     unspecified_special_event  typical_exercise_activity  \\\n",
+       "sample_idx time_idx                                                         \n",
+       "0          0                               NaN                        NaN   \n",
+       "           1                               NaN                        NaN   \n",
+       "           2                               NaN                        NaN   \n",
+       "           3                               NaN                        NaN   \n",
+       "           4                               NaN                        NaN   \n",
+       "...                                        ...                        ...   \n",
+       "69         146                             NaN                        NaN   \n",
+       "           147                             NaN                        NaN   \n",
+       "           148                             NaN                        NaN   \n",
+       "           149                             NaN                        NaN   \n",
+       "           150                             NaN                        NaN   \n",
+       "\n",
+       "                     less-than-usual_exercise_activity  \\\n",
+       "sample_idx time_idx                                      \n",
+       "0          0                                       NaN   \n",
+       "           1                                       NaN   \n",
+       "           2                                       NaN   \n",
+       "           3                                       NaN   \n",
+       "           4                                       NaN   \n",
+       "...                                                ...   \n",
+       "69         146                                     NaN   \n",
+       "           147                                     NaN   \n",
+       "           148                                     NaN   \n",
+       "           149                                     NaN   \n",
+       "           150                                     NaN   \n",
+       "\n",
+       "                     post-supper_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                          \n",
+       "0          0                                           NaN   \n",
+       "           1                                           NaN   \n",
+       "           2                                           NaN   \n",
+       "           3                                           NaN   \n",
+       "           4                                           NaN   \n",
+       "...                                                    ...   \n",
+       "69         146                                         NaN   \n",
+       "           147                                         NaN   \n",
+       "           148                                         NaN   \n",
+       "           149                                         NaN   \n",
+       "           150                                         NaN   \n",
+       "\n",
+       "                     more-than-usual_exercise_activity  \\\n",
+       "sample_idx time_idx                                      \n",
+       "0          0                                       NaN   \n",
+       "           1                                       NaN   \n",
+       "           2                                       NaN   \n",
+       "           3                                       NaN   \n",
+       "           4                                       NaN   \n",
+       "...                                                ...   \n",
+       "69         146                                     NaN   \n",
+       "           147                                     NaN   \n",
+       "           148                                     NaN   \n",
+       "           149                                     NaN   \n",
+       "           150                                     NaN   \n",
+       "\n",
+       "                     unspecified_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                          \n",
+       "0          0                                           NaN   \n",
+       "           1                                           NaN   \n",
+       "           2                                         123.0   \n",
+       "           3                                           NaN   \n",
+       "           4                                           NaN   \n",
+       "...                                                    ...   \n",
+       "69         146                                       145.0   \n",
+       "           147                                         NaN   \n",
+       "           148                                         NaN   \n",
+       "           149                                         NaN   \n",
+       "           150                                         NaN   \n",
+       "\n",
+       "                     pre-snack_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                        \n",
+       "0          0                                         NaN   \n",
+       "           1                                         NaN   \n",
+       "           2                                         NaN   \n",
+       "           3                                         NaN   \n",
+       "           4                                         NaN   \n",
+       "...                                                  ...   \n",
+       "69         146                                       NaN   \n",
+       "           147                                       NaN   \n",
+       "           148                                       NaN   \n",
+       "           149                                       NaN   \n",
+       "           150                                       NaN   \n",
+       "\n",
+       "                     ultralente_insulin_dose  less-than-usual_meal_ingestion  \\\n",
+       "sample_idx time_idx                                                            \n",
+       "0          0                             NaN                             NaN   \n",
+       "           1                             NaN                             NaN   \n",
+       "           2                             NaN                             NaN   \n",
+       "           3                             NaN                             NaN   \n",
+       "           4                             NaN                             NaN   \n",
+       "...                                      ...                             ...   \n",
+       "69         146                           NaN                             NaN   \n",
+       "           147                           NaN                             NaN   \n",
+       "           148                           NaN                             NaN   \n",
+       "           149                           NaN                             NaN   \n",
+       "           150                           NaN                             NaN   \n",
+       "\n",
+       "                     pre-supper_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                         \n",
+       "0          0                                          NaN   \n",
+       "           1                                        119.0   \n",
+       "           2                                          NaN   \n",
+       "           3                                          NaN   \n",
+       "           4                                          NaN   \n",
+       "...                                                   ...   \n",
+       "69         146                                        NaN   \n",
+       "           147                                        NaN   \n",
+       "           148                                        NaN   \n",
+       "           149                                        NaN   \n",
+       "           150                                        NaN   \n",
+       "\n",
+       "                     regular_insulin_dose  \\\n",
+       "sample_idx time_idx                         \n",
+       "0          0                          9.0   \n",
+       "           1                          7.0   \n",
+       "           2                          NaN   \n",
+       "           3                         10.0   \n",
+       "           4                          2.0   \n",
+       "...                                   ...   \n",
+       "69         146                        NaN   \n",
+       "           147                        1.0   \n",
+       "           148                        NaN   \n",
+       "           149                        NaN   \n",
+       "           150                        NaN   \n",
+       "\n",
+       "                     pre-breakfast_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                            \n",
+       "0          0                                           100.0   \n",
+       "           1                                             NaN   \n",
+       "           2                                             NaN   \n",
+       "           3                                           216.0   \n",
+       "           4                                             NaN   \n",
+       "...                                                      ...   \n",
+       "69         146                                           NaN   \n",
+       "           147                                         259.0   \n",
+       "           148                                           NaN   \n",
+       "           149                                           NaN   \n",
+       "           150                                           NaN   \n",
+       "\n",
+       "                     nph_insulin_dose  \\\n",
+       "sample_idx time_idx                     \n",
+       "0          0                     13.0   \n",
+       "           1                      NaN   \n",
+       "           2                      NaN   \n",
+       "           3                     13.0   \n",
+       "           4                      NaN   \n",
+       "...                               ...   \n",
+       "69         146                    NaN   \n",
+       "           147                    7.0   \n",
+       "           148                    7.0   \n",
+       "           149                    7.0   \n",
+       "           150                    7.0   \n",
+       "\n",
+       "                     post-breakfast_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                             \n",
+       "0          0                                              NaN   \n",
+       "           1                                              NaN   \n",
+       "           2                                              NaN   \n",
+       "           3                                              NaN   \n",
+       "           4                                              NaN   \n",
+       "...                                                       ...   \n",
+       "69         146                                            NaN   \n",
+       "           147                                            NaN   \n",
+       "           148                                            NaN   \n",
+       "           149                                            NaN   \n",
+       "           150                                            NaN   \n",
+       "\n",
+       "                     pre-lunch_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                        \n",
+       "0          0                                         NaN   \n",
+       "           1                                         NaN   \n",
+       "           2                                         NaN   \n",
+       "           3                                         NaN   \n",
+       "           4                                         NaN   \n",
+       "...                                                  ...   \n",
+       "69         146                                       NaN   \n",
+       "           147                                       NaN   \n",
+       "           148                                       NaN   \n",
+       "           149                                       NaN   \n",
+       "           150                                       NaN   \n",
+       "\n",
+       "                     typical_meal_ingestion  \n",
+       "sample_idx time_idx                          \n",
+       "0          0                            NaN  \n",
+       "           1                            NaN  \n",
+       "           2                            NaN  \n",
+       "           3                            NaN  \n",
+       "           4                            NaN  \n",
+       "...                                     ...  \n",
+       "69         146                          NaN  \n",
+       "           147                          NaN  \n",
+       "           148                          NaN  \n",
+       "           149                          NaN  \n",
+       "           150                          NaN  \n",
        "\n",
        "[18199 rows x 18 columns]"
       ]
@@ -2013,7 +2235,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "DummyTemporalPredictionDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "TemporalPredictionDataset(\n",
       "    time_series=TimeSeriesSamples([100, *, 5]),\n",
       "    static=StaticSamples([100, 3]),\n",
@@ -2449,7 +2677,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "PBCDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "TimeToEventAnalysisDataset(\n",
       "    time_series=TimeSeriesSamples([312, *, 14]),\n",
       "    static=StaticSamples([312, 1]),\n",
@@ -2715,19 +2949,47 @@
       ],
       "text/plain": [
        "TimeSeriesSamples with data:\n",
-       "                     drug  ascites  ...  prothrombin       age\n",
-       "sample_idx time_idx                 ...                       \n",
-       "1          0.569489   0.0      1.0  ...     0.136768  0.248058\n",
-       "           1.095170   0.0      1.0  ...     0.813132  0.248058\n",
-       "2          5.319790   0.0      1.0  ...     0.339677  1.292856\n",
-       "           6.261636   0.0      1.0  ...     0.339677  1.292856\n",
-       "           7.266455   0.0      1.0  ...     0.339677  1.292856\n",
-       "...                   ...      ...  ...          ...       ...\n",
-       "312        1.045888   1.0      0.0  ...     1.624769 -1.962482\n",
-       "           1.867265   1.0      0.0  ...     1.354223 -1.962482\n",
-       "           2.921367   1.0      0.0  ...     0.474950 -1.962482\n",
-       "           3.425145   1.0      0.0  ...    -0.066141 -1.962482\n",
-       "           3.989158   1.0      0.0  ...    -0.133778 -1.962482\n",
+       "                     drug  ascites  hepatomegaly  spiders  edema  histologic  \\\n",
+       "sample_idx time_idx                                                            \n",
+       "1          0.569489   0.0      1.0           1.0      1.0    1.0         3.0   \n",
+       "           1.095170   0.0      1.0           1.0      1.0    1.0         3.0   \n",
+       "2          5.319790   0.0      1.0           1.0      1.0    1.0         2.0   \n",
+       "           6.261636   0.0      1.0           1.0      1.0    1.0         2.0   \n",
+       "           7.266455   0.0      1.0           1.0      1.0    1.0         2.0   \n",
+       "...                   ...      ...           ...      ...    ...         ...   \n",
+       "312        1.045888   1.0      0.0           0.0      1.0    2.0         2.0   \n",
+       "           1.867265   1.0      0.0           0.0      1.0    2.0         1.0   \n",
+       "           2.921367   1.0      0.0           0.0      0.0    0.0         1.0   \n",
+       "           3.425145   1.0      0.0           0.0      0.0    0.0         1.0   \n",
+       "           3.989158   1.0      0.0           0.0      1.0    0.0         1.0   \n",
+       "\n",
+       "                     serBilir   serChol   albumin  alkaline      SGOT  \\\n",
+       "sample_idx time_idx                                                     \n",
+       "1          0.569489  3.281890  0.000000 -0.894575  0.195532 -1.485263   \n",
+       "           1.095170  2.015877 -0.469461 -1.570646  0.285613  0.195488   \n",
+       "2          5.319790  0.172710 -0.658914 -1.431455 -0.605844 -0.442126   \n",
+       "           6.261636 -0.013468 -0.603657 -1.172958 -0.512364 -0.046806   \n",
+       "           7.266455  0.098239  0.000000 -1.312149 -0.443529  0.293680   \n",
+       "...                       ...       ...       ...       ...       ...   \n",
+       "312        1.045888  3.672865  3.319599  0.059878  1.385274  0.986129   \n",
+       "           1.867265  2.350998  2.901224 -0.099197  0.916176  0.641817   \n",
+       "           2.921367  0.694010 -0.066873  0.338261  0.327254  0.552551   \n",
+       "           3.425145  0.340271  0.000000 -0.377580  0.251620  0.016956   \n",
+       "           3.989158  0.507832  2.017110  0.795603  0.622990  0.169983   \n",
+       "\n",
+       "                     platelets  prothrombin       age  \n",
+       "sample_idx time_idx                                    \n",
+       "1          0.569489  -0.529101     0.136768  0.248058  \n",
+       "           1.095170  -0.456022     0.813132  0.248058  \n",
+       "2          5.319790  -1.395605     0.339677  1.292856  \n",
+       "           6.261636  -1.259888     0.339677  1.292856  \n",
+       "           7.266455  -1.364286     0.339677  1.292856  \n",
+       "...                        ...          ...       ...  \n",
+       "312        1.045888  -1.103291     1.624769 -1.962482  \n",
+       "           1.867265  -0.998892     1.354223 -1.962482  \n",
+       "           2.921367  -0.894494     0.474950 -1.962482  \n",
+       "           3.425145  -0.466462    -0.066141 -1.962482  \n",
+       "           3.989158  -0.351624    -0.133778 -1.962482  \n",
        "\n",
        "[1945 rows x 14 columns]"
       ]
@@ -2958,7 +3220,13 @@
       "\n",
       "PKPDDataSource loads the following dataset:\n",
       "\n",
-      "Generating simple PKPD dataset with random seed 100...\n",
+      "Generating simple PKPD dataset with random seed 100...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "OneOffTreatmentEffectsDataset(\n",
       "    time_series=TimeSeriesSamples([40, *, 2]),\n",
       "    predictive=OneOffTreatmentEffectsTaskData(\n",
@@ -3458,7 +3726,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "DummyTemporalTreatmentEffectsDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "TemporalTreatmentEffectsDataset(\n",
       "    time_series=TimeSeriesSamples([100, *, 5]),\n",
       "    static=StaticSamples([100, 3]),\n",
diff --git a/docs/user_guide/extending/tutorial01_custom_method.ipynb b/docs/user_guide/extending/tutorial01_custom_method.ipynb
index da1e830f..e4cfe838 100644
--- a/docs/user_guide/extending/tutorial01_custom_method.ipynb
+++ b/docs/user_guide/extending/tutorial01_custom_method.ipynb
@@ -105,13 +105,7 @@
        "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'prediction'</span>: <span style=\"font-weight: bold\">{</span>\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span>\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span>\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'cde_classifier'</span>,\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'ode_classifier'</span>,\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'nn_classifier'</span>,\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'laplace_ode_classifier'</span>,\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'my_classifier'</span>\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"font-weight: bold\">]</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'cde_classifier'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'ode_classifier'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'nn_classifier'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'laplace_ode_classifier'</span><span style=\"font-weight: bold\">]</span>,\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'laplace_ode_regressor'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'nn_regressor'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'ode_regressor'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'cde_regressor'</span><span style=\"font-weight: bold\">]</span>\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>,\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'temporal'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'seq2seq_classifier'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'seq2seq_regressor'</span><span style=\"font-weight: bold\">]}</span>\n",
@@ -140,13 +134,7 @@
        "\u001b[1m{\u001b[0m\n",
        "\u001b[2;32m│   \u001b[0m\u001b[32m'prediction'\u001b[0m: \u001b[1m{\u001b[0m\n",
        "\u001b[2;32m│   │   \u001b[0m\u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\n",
-       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'cde_classifier'\u001b[0m,\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'ode_classifier'\u001b[0m,\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'nn_classifier'\u001b[0m,\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'laplace_ode_classifier'\u001b[0m,\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'my_classifier'\u001b[0m\n",
-       "\u001b[2;32m│   │   │   \u001b[0m\u001b[1m]\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'cde_classifier'\u001b[0m, \u001b[32m'ode_classifier'\u001b[0m, \u001b[32m'nn_classifier'\u001b[0m, \u001b[32m'laplace_ode_classifier'\u001b[0m\u001b[1m]\u001b[0m,\n",
        "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'laplace_ode_regressor'\u001b[0m, \u001b[32m'nn_regressor'\u001b[0m, \u001b[32m'ode_regressor'\u001b[0m, \u001b[32m'cde_regressor'\u001b[0m\u001b[1m]\u001b[0m\n",
        "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m,\n",
        "\u001b[2;32m│   │   \u001b[0m\u001b[32m'temporal'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'seq2seq_classifier'\u001b[0m\u001b[1m]\u001b[0m, \u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'seq2seq_regressor'\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\n",
diff --git a/docs/user_guide/extending/tutorial03_custom_datasource.ipynb b/docs/user_guide/extending/tutorial03_custom_datasource.ipynb
index 4a2efe63..79501854 100644
--- a/docs/user_guide/extending/tutorial03_custom_datasource.ipynb
+++ b/docs/user_guide/extending/tutorial03_custom_datasource.ipynb
@@ -73,7 +73,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Remember you can also see the existing method plugins and how they correspond to different categories, as follows:"
+    "Remember you can also see the existing data source plugins and how they correspond to different categories, as follows:"
    ]
   },
   {
@@ -348,7 +348,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "<__main__.MyDataSource object at 0x7fa2cf032850>\n"
+      "<__main__.MyDataSource object at 0x7f5f36a6cbb0>\n"
      ]
     }
    ],
diff --git a/docs/user_guide/extending/tutorial04_custom_metric.ipynb b/docs/user_guide/extending/tutorial04_custom_metric.ipynb
new file mode 100644
index 00000000..ac95b2f6
--- /dev/null
+++ b/docs/user_guide/extending/tutorial04_custom_metric.ipynb
@@ -0,0 +1,471 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "",
+   "metadata": {},
+   "source": [
+    "[![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial04_custom_metric.ipynb)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Extending TemporAI Tutorial 04: Writing a Custom Metric Plugin\n",
+    "\n",
+    "This tutorial shows how to extend TemporAI by wring a custom *metric* plugin."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Note**\n",
+    "\n",
+    "See also \"Writing a Custom `Plugin` 101\" section in \"Writing a Custom Method Plugin\" tutorial."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Inherit from the appropriate **base class** for the category of the metric plugin you are writing.\n",
+    "\n",
+    "You need to find which category of metric plugin you are writing.\n",
+    "\n",
+    "You can view all the different metric plugin categories as so:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['prediction.one_off.classification',\n",
+       " 'prediction.one_off.regression',\n",
+       " 'time_to_event']"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from tempor import plugin_loader\n",
+    "\n",
+    "plugin_categories = plugin_loader.list_categories(plugin_type=\"metric\")\n",
+    "\n",
+    "list(plugin_categories.keys())"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Remember you can also see the existing metric plugins and how they correspond to different categories, as follows:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'prediction'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'mse'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'mae'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'r2'</span><span style=\"font-weight: bold\">]</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'accuracy'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'kappa'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'kappa_quadratic'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'mcc'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'aucprc'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'aucroc'</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"font-weight: bold\">]</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"font-weight: bold\">}</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'time_to_event'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'c_index'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'brier_score'</span><span style=\"font-weight: bold\">]</span>\n",
+       "<span style=\"font-weight: bold\">}</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[32m'prediction'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   │   \u001b[0m\u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'mse'\u001b[0m, \u001b[32m'mae'\u001b[0m, \u001b[32m'r2'\u001b[0m\u001b[1m]\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'accuracy'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'kappa'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'kappa_quadratic'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'mcc'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'aucprc'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'aucroc'\u001b[0m\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[1m]\u001b[0m\n",
+       "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[1m}\u001b[0m,\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[32m'time_to_event'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'c_index'\u001b[0m, \u001b[32m'brier_score'\u001b[0m\u001b[1m]\u001b[0m\n",
+       "\u001b[1m}\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "all_plugins = plugin_loader.list(plugin_type=\"metric\")\n",
+    "\n",
+    "from rich.pretty import pprint  # For prettifying the print output only.\n",
+    "\n",
+    "pprint(all_plugins, indent_guides=True)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's say you would like to write a plugin of category `\"prediction.one_off.regression\"`.\n",
+    "\n",
+    "You can find which base class you need to inherit from as follows."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Base classes for all categories:\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
+       "    <span style=\"color: #008000; text-decoration-color: #008000\">'prediction.one_off.classification'</span>: <span style=\"font-weight: bold\">&lt;</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff; font-weight: bold\">class</span><span style=\"color: #000000; text-decoration-color: #000000\"> </span><span style=\"color: #008000; text-decoration-color: #008000\">'tempor.metrics.metric.OneOffClassificationMetric'</span><span style=\"color: #000000; text-decoration-color: #000000\">&gt;,</span>\n",
+       "<span style=\"color: #000000; text-decoration-color: #000000\">    </span><span style=\"color: #008000; text-decoration-color: #008000\">'prediction.one_off.regression'</span><span style=\"color: #000000; text-decoration-color: #000000\">: &lt;class </span><span style=\"color: #008000; text-decoration-color: #008000\">'tempor.metrics.metric.OneOffRegressionMetric'</span><span style=\"color: #000000; text-decoration-color: #000000\">&gt;,</span>\n",
+       "<span style=\"color: #000000; text-decoration-color: #000000\">    </span><span style=\"color: #008000; text-decoration-color: #008000\">'time_to_event'</span><span style=\"color: #000000; text-decoration-color: #000000\">: &lt;class </span><span style=\"color: #008000; text-decoration-color: #008000\">'tempor.metrics.metric.TimeToEventMetric'</span><span style=\"font-weight: bold\">&gt;</span>\n",
+       "<span style=\"font-weight: bold\">}</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[1m{\u001b[0m\n",
+       "    \u001b[32m'prediction.one_off.classification'\u001b[0m: \u001b[1m<\u001b[0m\u001b[1;95mclass\u001b[0m\u001b[39m \u001b[0m\u001b[32m'tempor.metrics.metric.OneOffClassificationMetric'\u001b[0m\u001b[39m>,\u001b[0m\n",
+       "\u001b[39m    \u001b[0m\u001b[32m'prediction.one_off.regression'\u001b[0m\u001b[39m: <class \u001b[0m\u001b[32m'tempor.metrics.metric.OneOffRegressionMetric'\u001b[0m\u001b[39m>,\u001b[0m\n",
+       "\u001b[39m    \u001b[0m\u001b[32m'time_to_event'\u001b[0m\u001b[39m: <class \u001b[0m\u001b[32m'tempor.metrics.metric.TimeToEventMetric'\u001b[0m\u001b[1m>\u001b[0m\n",
+       "\u001b[1m}\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Base class you need:\n",
+      "<class 'tempor.metrics.metric.OneOffRegressionMetric'>\n"
+     ]
+    }
+   ],
+   "source": [
+    "plugin_categories = plugin_loader.list_categories(plugin_type=\"metric\")\n",
+    "\n",
+    "print(\"Base classes for all categories:\")\n",
+    "pprint(plugin_categories, indent_guides=False)\n",
+    "\n",
+    "print(\"Base class you need:\")\n",
+    "print(plugin_categories[\"prediction.one_off.regression\"])"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can then find the class in the TemporAI source code, to see its method signatures etc."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Implement the **methods** the plugin needs.\n",
+    "\n",
+    "`Metric` plugins require the following methods to be implemented:\n",
+    "* `direction` property, which returns either `\"minimize\"` or `\"maximize\"`, representing the \"good\" direction of the metric.\n",
+    "* `_evaluate()` which takes in the actual and predicted values and returns the evaluated metric(s).\n",
+    "\n",
+    "If you haven't implemented some required method for the plugin, Python will notify you by raising an exception when you\n",
+    "attempt to instantiate your plugin (see [Python `abc`](https://docs.python.org/3/library/abc.html)).\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Register** the plugin with TemporAI.\n",
+    "\n",
+    "Registering your plugin with TemporAI is very simple, you need to use the `register_plugin` decorator,\n",
+    "as shown in the example below.\n",
+    "\n",
+    "You will need to specify the `name` of your plugin and its `category` in the decorator.\n",
+    "\n",
+    "The `plugin_type` needs to be set to `\"datasource\"`.\n",
+    "\n",
+    "```python\n",
+    "from tempor.core.plugins import register_plugin\n",
+    "\n",
+    "@register_plugin(name=\"my_metric\", category=\"prediction.one_off.regression\", plugin_type=\"metric\")\n",
+    "class MyMetric(OneOffRegressionMetric):\n",
+    "    ...\n",
+    "```"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Example\n",
+    "\n",
+    "Now putting this together in a minimal example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Any\n",
+    "\n",
+    "import numpy as np\n",
+    "from sklearn.metrics import d2_tweedie_score\n",
+    "\n",
+    "from tempor.core import plugins\n",
+    "from tempor.metrics import metric, metric_typing\n",
+    "\n",
+    "\n",
+    "@plugins.register_plugin(name=\"my_metric\", category=\"prediction.one_off.regression\", plugin_type=\"metric\")\n",
+    "class MyOneOffRegressionMetric(metric.OneOffRegressionMetric):\n",
+    "    \"\"\"My custom metric, here we use\n",
+    "    D^2 regression score function, fraction of Tweedie deviance explained,\n",
+    "    as per `sklearn`.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    @property\n",
+    "    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102\n",
+    "        return \"maximize\"\n",
+    "\n",
+    "    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:\n",
+    "        return d2_tweedie_score(actual, predicted, power=kwargs.get(\"power\", 0))"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We now see our plugin in TemporAI:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'prediction'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'mse'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'mae'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'r2'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'my_metric'</span><span style=\"font-weight: bold\">]</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'accuracy'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'kappa'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'kappa_quadratic'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'mcc'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'aucprc'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'aucroc'</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"font-weight: bold\">]</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"font-weight: bold\">}</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'time_to_event'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'c_index'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'brier_score'</span><span style=\"font-weight: bold\">]</span>\n",
+       "<span style=\"font-weight: bold\">}</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[32m'prediction'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   │   \u001b[0m\u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'mse'\u001b[0m, \u001b[32m'mae'\u001b[0m, \u001b[32m'r2'\u001b[0m, \u001b[32m'my_metric'\u001b[0m\u001b[1m]\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'accuracy'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'kappa'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'kappa_quadratic'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'mcc'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'aucprc'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'aucroc'\u001b[0m\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[1m]\u001b[0m\n",
+       "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[1m}\u001b[0m,\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[32m'time_to_event'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'c_index'\u001b[0m, \u001b[32m'brier_score'\u001b[0m\u001b[1m]\u001b[0m\n",
+       "\u001b[1m}\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "`my_metric` plugin found in the category 'prediction.one_off.regression': True\n"
+     ]
+    }
+   ],
+   "source": [
+    "from tempor import plugin_loader\n",
+    "\n",
+    "all_plugins = plugin_loader.list(plugin_type=\"metric\")\n",
+    "\n",
+    "pprint(all_plugins, indent_guides=True)\n",
+    "\n",
+    "my_metric_found = \"my_metric\" in all_plugins[\"prediction\"][\"one_off\"][\"regression\"]\n",
+    "print(f\"`my_metric` plugin found in the category 'prediction.one_off.regression': {my_metric_found}\")\n",
+    "assert my_metric_found"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The plugin can be used as normal."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "MyOneOffRegressionMetric(\n",
+      "    name='my_metric',\n",
+      "    description='My custom metric, here we use\\n    D^2 regression score function, fraction of Tweedie deviance explai...'\n",
+      ")\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Get the plugin.\n",
+    "\n",
+    "my_metric = plugin_loader.get(\"prediction.one_off.regression.my_metric\", plugin_type=\"metric\")\n",
+    "\n",
+    "print(my_metric)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.7492656592741931\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Use the metric.\n",
+    "\n",
+    "y_true = np.asarray([1, 2, 3])\n",
+    "y_pred = np.asarray([1, 3, 3])\n",
+    "\n",
+    "metric_value = my_metric(y_true, y_pred, power=2)\n",
+    "print(metric_value)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "RE_py38_temporai",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/docs/user_guide/usage/tutorial01_plugins.ipynb b/docs/user_guide/usage/tutorial01_plugins.ipynb
index 116bdf11..9bda2121 100644
--- a/docs/user_guide/usage/tutorial01_plugins.ipynb
+++ b/docs/user_guide/usage/tutorial01_plugins.ipynb
@@ -340,7 +340,7 @@
     {
      "data": {
       "text/plain": [
-       "['datasource', 'method']"
+       "['datasource', 'method', 'metric']"
       ]
      },
      "execution_count": null,
@@ -398,6 +398,31 @@
        "            <span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'synctwin_regressor'</span><span style=\"font-weight: bold\">]}</span>,\n",
        "            <span style=\"color: #008000; text-decoration-color: #008000\">'temporal'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'crn_classifier'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'crn_regressor'</span><span style=\"font-weight: bold\">]}</span>\n",
        "        <span style=\"font-weight: bold\">}</span>\n",
+       "    <span style=\"font-weight: bold\">}</span>,\n",
+       "    <span style=\"color: #008000; text-decoration-color: #008000\">'metric'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'prediction'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "            <span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "                <span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'mse'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'mae'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'r2'</span><span style=\"font-weight: bold\">]</span>,\n",
+       "                <span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span>\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'accuracy'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_micro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_macro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_weighted'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'kappa'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'kappa_quadratic'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'recall_micro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'recall_macro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'recall_weighted'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'precision_micro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'precision_macro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'precision_weighted'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'mcc'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'aucprc'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'aucroc'</span>\n",
+       "                <span style=\"font-weight: bold\">]</span>\n",
+       "            <span style=\"font-weight: bold\">}</span>\n",
+       "        <span style=\"font-weight: bold\">}</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'time_to_event'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'c_index'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'brier_score'</span><span style=\"font-weight: bold\">]</span>\n",
        "    <span style=\"font-weight: bold\">}</span>\n",
        "<span style=\"font-weight: bold\">}</span>\n",
        "</pre>\n"
@@ -434,6 +459,31 @@
        "            \u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'synctwin_regressor'\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m,\n",
        "            \u001b[32m'temporal'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'crn_classifier'\u001b[0m\u001b[1m]\u001b[0m, \u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'crn_regressor'\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\n",
        "        \u001b[1m}\u001b[0m\n",
+       "    \u001b[1m}\u001b[0m,\n",
+       "    \u001b[32m'metric'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "        \u001b[32m'prediction'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "            \u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "                \u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'mse'\u001b[0m, \u001b[32m'mae'\u001b[0m, \u001b[32m'r2'\u001b[0m\u001b[1m]\u001b[0m,\n",
+       "                \u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\n",
+       "                    \u001b[32m'accuracy'\u001b[0m,\n",
+       "                    \u001b[32m'f1_score_micro'\u001b[0m,\n",
+       "                    \u001b[32m'f1_score_macro'\u001b[0m,\n",
+       "                    \u001b[32m'f1_score_weighted'\u001b[0m,\n",
+       "                    \u001b[32m'kappa'\u001b[0m,\n",
+       "                    \u001b[32m'kappa_quadratic'\u001b[0m,\n",
+       "                    \u001b[32m'recall_micro'\u001b[0m,\n",
+       "                    \u001b[32m'recall_macro'\u001b[0m,\n",
+       "                    \u001b[32m'recall_weighted'\u001b[0m,\n",
+       "                    \u001b[32m'precision_micro'\u001b[0m,\n",
+       "                    \u001b[32m'precision_macro'\u001b[0m,\n",
+       "                    \u001b[32m'precision_weighted'\u001b[0m,\n",
+       "                    \u001b[32m'mcc'\u001b[0m,\n",
+       "                    \u001b[32m'aucprc'\u001b[0m,\n",
+       "                    \u001b[32m'aucroc'\u001b[0m\n",
+       "                \u001b[1m]\u001b[0m\n",
+       "            \u001b[1m}\u001b[0m\n",
+       "        \u001b[1m}\u001b[0m,\n",
+       "        \u001b[32m'time_to_event'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'c_index'\u001b[0m, \u001b[32m'brier_score'\u001b[0m\u001b[1m]\u001b[0m\n",
        "    \u001b[1m}\u001b[0m\n",
        "\u001b[1m}\u001b[0m\n"
       ]
diff --git a/docs/user_guide/usage/tutorial02_imputation.ipynb b/docs/user_guide/usage/tutorial02_imputation.ipynb
index 8d3ece58..a4955c6a 100644
--- a/docs/user_guide/usage/tutorial02_imputation.ipynb
+++ b/docs/user_guide/usage/tutorial02_imputation.ipynb
@@ -84,7 +84,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
      ]
     },
     {
@@ -289,41 +289,245 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.22005514968324613 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21670750510884584 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.2166465658117811 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.2166456843686172 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 7\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 8\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 9\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 10\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.22005514968324613 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21670750510884584 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.2166465658117811 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.2166456843686172 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 7\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 8\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 9\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 10\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
      ]
     },
     {
diff --git a/docs/user_guide/usage/tutorial04_prediction.ipynb b/docs/user_guide/usage/tutorial04_prediction.ipynb
index d653ba1b..50b85cf7 100644
--- a/docs/user_guide/usage/tutorial04_prediction.ipynb
+++ b/docs/user_guide/usage/tutorial04_prediction.ipynb
@@ -271,11 +271,35 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.68824702501297, validation loss: 0.6883962154388428\n",
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:10| train loss: 0.6999298334121704, validation loss: 0.6879696846008301\n",
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:20| train loss: 0.6818332672119141, validation loss: 0.6867876648902893\n",
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:30| train loss: 0.6943979263305664, validation loss: 0.685184121131897\n",
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:40| train loss: 0.6284184455871582, validation loss: 0.6934055089950562\n"
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.68824702501297, validation loss: 0.6883962154388428\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:10| train loss: 0.6999298930168152, validation loss: 0.6879696249961853\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:20| train loss: 0.6818333268165588, validation loss: 0.6867876052856445\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:30| train loss: 0.694374680519104, validation loss: 0.6851812601089478\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:40| train loss: 0.6282318234443665, validation loss: 0.6936542391777039\n"
      ]
     },
     {
@@ -643,30 +667,114 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Preparing data for decoder training...\n",
-      "Preparing data for decoder training DONE.\n",
+      "Preparing data for decoder training...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Preparing data for decoder training DONE.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "=== Training stage: 1. Train encoder ===\n",
       "Epoch: 0, Loss: 68.64124755859375\n",
-      "Epoch: 1, Loss: 35.37634651184082\n",
-      "Epoch: 2, Loss: 21.269977149963378\n",
-      "Epoch: 3, Loss: 17.418054847717286\n",
-      "Epoch: 4, Loss: 13.851956520080567\n",
-      "Epoch: 5, Loss: 8.335217247009277\n",
-      "Epoch: 6, Loss: 6.385376853942871\n",
-      "Epoch: 7, Loss: 5.112505340576172\n",
-      "Epoch: 8, Loss: 4.690582332611084\n",
-      "Epoch: 9, Loss: 4.158801422119141\n",
-      "=== Training stage: 2. Train decoder ===\n",
-      "Epoch: 0, Loss: 29.110631188485222\n",
-      "Epoch: 1, Loss: 4.417614046603097\n",
-      "Epoch: 2, Loss: 3.908697671333463\n",
-      "Epoch: 3, Loss: 3.8485680392058788\n",
-      "Epoch: 4, Loss: 3.8276885692068126\n",
-      "Epoch: 5, Loss: 3.793576312405029\n",
-      "Epoch: 6, Loss: 3.749186389072364\n",
-      "Epoch: 7, Loss: 3.787092407949102\n",
-      "Epoch: 8, Loss: 3.7378030509586164\n",
-      "Epoch: 9, Loss: 3.717736062595655\n"
+      "Epoch: 1, Loss: 35.376345291137696\n",
+      "Epoch: 2, Loss: 21.269977149963378\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 3, Loss: 17.418056030273437\n",
+      "Epoch: 4, Loss: 13.851958045959472\n",
+      "Epoch: 5, Loss: 8.335217628479004\n",
+      "Epoch: 6, Loss: 6.385376873016358\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 7, Loss: 5.1125048828125\n",
+      "Epoch: 8, Loss: 4.690582141876221\n",
+      "Epoch: 9, Loss: 4.15880220413208\n",
+      "=== Training stage: 2. Train decoder ===\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 0, Loss: 29.110622051453316\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 1, Loss: 4.4235033275153635\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 2, Loss: 3.954428958957566\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 3, Loss: 3.8989672662446226\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 4, Loss: 3.8317033024119103\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 5, Loss: 3.8341090269664964\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 6, Loss: 3.717077842797864\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 7, Loss: 3.7477998285170564\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 8, Loss: 3.859872080239\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 9, Loss: 3.8217773019903305\n"
      ]
     }
    ],
@@ -716,28 +824,28 @@
        "    <tr>\n",
        "      <th rowspan=\"5\" valign=\"top\">0</th>\n",
        "      <th>11</th>\n",
-       "      <td>9.828533</td>\n",
-       "      <td>8.405809</td>\n",
+       "      <td>10.618200</td>\n",
+       "      <td>8.973257</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>12</th>\n",
-       "      <td>9.831676</td>\n",
-       "      <td>8.414574</td>\n",
+       "      <td>10.728621</td>\n",
+       "      <td>9.108867</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>13</th>\n",
-       "      <td>9.813161</td>\n",
-       "      <td>8.395016</td>\n",
+       "      <td>10.735380</td>\n",
+       "      <td>9.118745</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>14</th>\n",
-       "      <td>9.798813</td>\n",
-       "      <td>8.378744</td>\n",
+       "      <td>10.728010</td>\n",
+       "      <td>9.093669</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>15</th>\n",
-       "      <td>9.772892</td>\n",
-       "      <td>8.352385</td>\n",
+       "      <td>10.738855</td>\n",
+       "      <td>9.097306</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>...</th>\n",
@@ -748,28 +856,28 @@
        "    <tr>\n",
        "      <th rowspan=\"5\" valign=\"top\">99</th>\n",
        "      <th>17</th>\n",
-       "      <td>11.214181</td>\n",
-       "      <td>9.754721</td>\n",
+       "      <td>12.021331</td>\n",
+       "      <td>10.417010</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>18</th>\n",
-       "      <td>11.275115</td>\n",
-       "      <td>9.821613</td>\n",
+       "      <td>12.091549</td>\n",
+       "      <td>10.484824</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>19</th>\n",
-       "      <td>11.298031</td>\n",
-       "      <td>9.848743</td>\n",
+       "      <td>12.144796</td>\n",
+       "      <td>10.523143</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>20</th>\n",
-       "      <td>11.280746</td>\n",
-       "      <td>9.832634</td>\n",
+       "      <td>12.133350</td>\n",
+       "      <td>10.510198</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>21</th>\n",
-       "      <td>11.256371</td>\n",
-       "      <td>9.801851</td>\n",
+       "      <td>12.097664</td>\n",
+       "      <td>10.484853</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -778,19 +886,19 @@
       ],
       "text/plain": [
        "TimeSeriesSamples with data:\n",
-       "                             0         1\n",
-       "sample_idx time_idx                     \n",
-       "0          11         9.828533  8.405809\n",
-       "           12         9.831676  8.414574\n",
-       "           13         9.813161  8.395016\n",
-       "           14         9.798813  8.378744\n",
-       "           15         9.772892  8.352385\n",
-       "...                        ...       ...\n",
-       "99         17        11.214181  9.754721\n",
-       "           18        11.275115  9.821613\n",
-       "           19        11.298031  9.848743\n",
-       "           20        11.280746  9.832634\n",
-       "           21        11.256371  9.801851\n",
+       "                             0          1\n",
+       "sample_idx time_idx                      \n",
+       "0          11        10.618200   8.973257\n",
+       "           12        10.728621   9.108867\n",
+       "           13        10.735380   9.118745\n",
+       "           14        10.728010   9.093669\n",
+       "           15        10.738855   9.097306\n",
+       "...                        ...        ...\n",
+       "99         17        12.021331  10.417010\n",
+       "           18        12.091549  10.484824\n",
+       "           19        12.144796  10.523143\n",
+       "           20        12.133350  10.510198\n",
+       "           21        12.097664  10.484853\n",
        "\n",
        "[1000 rows x 2 columns]"
       ]
diff --git a/docs/user_guide/usage/tutorial05_time_to_event.ipynb b/docs/user_guide/usage/tutorial05_time_to_event.ipynb
index 9eead21d..cf34ad83 100644
--- a/docs/user_guide/usage/tutorial05_time_to_event.ipynb
+++ b/docs/user_guide/usage/tutorial05_time_to_event.ipynb
@@ -321,24 +321,24 @@
        "    <tr>\n",
        "      <th rowspan=\"3\" valign=\"top\">1</th>\n",
        "      <th>0.25</th>\n",
-       "      <td>0.422643</td>\n",
+       "      <td>0.422638</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.50</th>\n",
-       "      <td>0.626620</td>\n",
+       "      <td>0.626617</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.75</th>\n",
-       "      <td>0.717206</td>\n",
+       "      <td>0.717204</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th rowspan=\"2\" valign=\"top\">2</th>\n",
        "      <th>0.25</th>\n",
-       "      <td>0.221314</td>\n",
+       "      <td>0.221283</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.50</th>\n",
-       "      <td>0.288492</td>\n",
+       "      <td>0.288457</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>...</th>\n",
@@ -348,24 +348,24 @@
        "    <tr>\n",
        "      <th rowspan=\"2\" valign=\"top\">311</th>\n",
        "      <th>0.50</th>\n",
-       "      <td>0.039302</td>\n",
+       "      <td>0.039307</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.75</th>\n",
-       "      <td>0.048346</td>\n",
+       "      <td>0.048353</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th rowspan=\"3\" valign=\"top\">312</th>\n",
        "      <th>0.25</th>\n",
-       "      <td>0.248834</td>\n",
+       "      <td>0.248848</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.50</th>\n",
-       "      <td>0.390042</td>\n",
+       "      <td>0.390068</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.75</th>\n",
-       "      <td>0.444533</td>\n",
+       "      <td>0.444567</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -376,17 +376,17 @@
        "TimeSeriesSamples with data:\n",
        "                     risk_score\n",
        "sample_idx time_idx            \n",
-       "1          0.25        0.422643\n",
-       "           0.50        0.626620\n",
-       "           0.75        0.717206\n",
-       "2          0.25        0.221314\n",
-       "           0.50        0.288492\n",
+       "1          0.25        0.422638\n",
+       "           0.50        0.626617\n",
+       "           0.75        0.717204\n",
+       "2          0.25        0.221283\n",
+       "           0.50        0.288457\n",
        "...                         ...\n",
-       "311        0.50        0.039302\n",
-       "           0.75        0.048346\n",
-       "312        0.25        0.248834\n",
-       "           0.50        0.390042\n",
-       "           0.75        0.444533\n",
+       "311        0.50        0.039307\n",
+       "           0.75        0.048353\n",
+       "312        0.25        0.248848\n",
+       "           0.50        0.390068\n",
+       "           0.75        0.444567\n",
        "\n",
        "[936 rows x 1 columns]"
       ]
diff --git a/docs/user_guide/usage/tutorial06_treatments.ipynb b/docs/user_guide/usage/tutorial06_treatments.ipynb
index 75f638f5..1f37d536 100644
--- a/docs/user_guide/usage/tutorial06_treatments.ipynb
+++ b/docs/user_guide/usage/tutorial06_treatments.ipynb
@@ -418,50 +418,218 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Preparing data for decoder training...\n",
-      "Preparing data for decoder training DONE.\n",
+      "Preparing data for decoder training...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Preparing data for decoder training DONE.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "=== Training stage: 1. Train encoder ===\n",
-      "Epoch: 0, Prediction Loss: 75.327, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 75.327\n",
-      "Epoch: 1, Prediction Loss: 37.270, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 37.270\n",
-      "Epoch: 2, Prediction Loss: 21.899, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 21.899\n",
-      "Epoch: 3, Prediction Loss: 18.649, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 18.649\n",
-      "Epoch: 4, Prediction Loss: 18.343, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 18.343\n",
-      "Epoch: 5, Prediction Loss: 15.802, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 15.802\n",
-      "Epoch: 6, Prediction Loss: 10.720, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 10.720\n",
-      "Epoch: 7, Prediction Loss: 8.111, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 8.111\n",
-      "Epoch: 8, Prediction Loss: 6.840, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 6.840\n",
-      "Epoch: 9, Prediction Loss: 5.708, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.708\n",
-      "Epoch: 10, Prediction Loss: 5.061, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.061\n",
-      "Epoch: 11, Prediction Loss: 4.648, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.648\n",
-      "Epoch: 12, Prediction Loss: 4.420, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.420\n",
-      "Epoch: 13, Prediction Loss: 4.356, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.356\n",
-      "Epoch: 14, Prediction Loss: 4.175, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.175\n",
-      "Epoch: 15, Prediction Loss: 4.014, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.014\n",
-      "Epoch: 16, Prediction Loss: 3.899, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.899\n",
-      "Epoch: 17, Prediction Loss: 3.933, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.933\n",
-      "Epoch: 18, Prediction Loss: 3.940, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.940\n",
-      "Epoch: 19, Prediction Loss: 3.844, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.844\n",
-      "=== Training stage: 2. Train decoder ===\n",
-      "Epoch: 0, Prediction Loss: 34.517, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 34.517\n",
-      "Epoch: 1, Prediction Loss: 7.175, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 7.175\n",
-      "Epoch: 2, Prediction Loss: 3.985, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.985\n",
-      "Epoch: 3, Prediction Loss: 3.831, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.831\n",
-      "Epoch: 4, Prediction Loss: 3.851, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.851\n",
-      "Epoch: 5, Prediction Loss: 3.730, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.730\n",
-      "Epoch: 6, Prediction Loss: 3.767, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.767\n",
-      "Epoch: 7, Prediction Loss: 3.737, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.737\n",
-      "Epoch: 8, Prediction Loss: 3.678, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.678\n",
-      "Epoch: 9, Prediction Loss: 3.659, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.659\n",
-      "Epoch: 10, Prediction Loss: 3.698, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.698\n",
-      "Epoch: 11, Prediction Loss: 3.684, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.684\n",
-      "Epoch: 12, Prediction Loss: 3.678, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.678\n",
-      "Epoch: 13, Prediction Loss: 3.649, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.649\n",
-      "Epoch: 14, Prediction Loss: 3.617, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.617\n",
-      "Epoch: 15, Prediction Loss: 3.670, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.670\n",
-      "Epoch: 16, Prediction Loss: 3.638, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.638\n",
-      "Epoch: 17, Prediction Loss: 3.655, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.655\n",
-      "Epoch: 18, Prediction Loss: 3.654, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.654\n",
-      "Epoch: 19, Prediction Loss: 3.668, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.668\n"
+      "Epoch: 0, Prediction Loss: 75.212, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 75.212\n",
+      "Epoch: 1, Prediction Loss: 32.453, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 32.453\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 2, Prediction Loss: 19.389, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 19.389\n",
+      "Epoch: 3, Prediction Loss: 19.679, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 19.679\n",
+      "Epoch: 4, Prediction Loss: 19.588, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 19.588\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 5, Prediction Loss: 14.915, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 14.915\n",
+      "Epoch: 6, Prediction Loss: 10.608, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 10.608\n",
+      "Epoch: 7, Prediction Loss: 8.684, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 8.684\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 8, Prediction Loss: 6.953, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 6.953\n",
+      "Epoch: 9, Prediction Loss: 5.645, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.645\n",
+      "Epoch: 10, Prediction Loss: 5.060, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.060\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 11, Prediction Loss: 4.658, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.658\n",
+      "Epoch: 12, Prediction Loss: 4.398, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.398\n",
+      "Epoch: 13, Prediction Loss: 4.265, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.265\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 14, Prediction Loss: 4.170, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.170\n",
+      "Epoch: 15, Prediction Loss: 4.054, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.054\n",
+      "Epoch: 16, Prediction Loss: 3.949, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.949\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 17, Prediction Loss: 3.940, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.940\n",
+      "Epoch: 18, Prediction Loss: 3.918, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.918\n",
+      "Epoch: 19, Prediction Loss: 3.851, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.851\n",
+      "=== Training stage: 2. Train decoder ===\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 0, Prediction Loss: 34.622, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 34.622\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 1, Prediction Loss: 5.329, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.329\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 2, Prediction Loss: 3.826, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.826\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 3, Prediction Loss: 3.767, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.767\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 4, Prediction Loss: 3.746, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.746\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 5, Prediction Loss: 3.741, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.741\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 6, Prediction Loss: 3.726, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.726\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 7, Prediction Loss: 3.694, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.694\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 8, Prediction Loss: 3.723, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.723\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 9, Prediction Loss: 3.730, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.730\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 10, Prediction Loss: 3.670, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.670\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 11, Prediction Loss: 3.691, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.691\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 12, Prediction Loss: 3.726, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.726\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 13, Prediction Loss: 3.701, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.701\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 14, Prediction Loss: 3.696, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.696\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 15, Prediction Loss: 3.668, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.668\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 16, Prediction Loss: 3.612, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.612\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 17, Prediction Loss: 3.718, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.718\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 18, Prediction Loss: 3.717, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.717\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 19, Prediction Loss: 3.658, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.658\n"
      ]
     }
    ],
@@ -491,23 +659,23 @@
       "TimeSeries() with data:\n",
       "                 0         1\n",
       "time_idx                    \n",
-      "5         5.993054  4.916986\n",
-      "6         6.041887  4.971257\n",
-      "7         6.043185  4.972708\n",
-      "8         6.043219  4.972745\n",
-      "9         6.043220  4.972746\n",
-      "10        6.043220  4.972746\n",
+      "5         5.966007  4.997900\n",
+      "6         6.120805  5.138746\n",
+      "7         6.138720  5.155504\n",
+      "8         6.140833  5.157483\n",
+      "9         6.141082  5.157716\n",
+      "10        6.141112  5.157744\n",
       "\n",
       "Treatment scenario 1, [0 0 0 0 0 0]\n",
       "TimeSeries() with data:\n",
       "                 0         1\n",
       "time_idx                    \n",
-      "5         6.361402  4.933103\n",
-      "6         6.410686  4.987370\n",
-      "7         6.412068  4.988914\n",
-      "8         6.412107  4.988956\n",
-      "9         6.412108  4.988957\n",
-      "10        6.412108  4.988957\n",
+      "5         6.303439  4.996541\n",
+      "6         6.467038  5.143897\n",
+      "7         6.486365  5.162077\n",
+      "8         6.488708  5.164285\n",
+      "9         6.488993  5.164552\n",
+      "10        6.489028  5.164585\n",
       "\n"
      ]
     }
diff --git a/docs/user_guide/usage/tutorial07_pipeline.ipynb b/docs/user_guide/usage/tutorial07_pipeline.ipynb
index 442eaab7..e467be77 100644
--- a/docs/user_guide/usage/tutorial07_pipeline.ipynb
+++ b/docs/user_guide/usage/tutorial07_pipeline.ipynb
@@ -116,8 +116,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:55:13 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n",
-      "2023-10-10 16:55:13 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
+      "2023-11-16 22:18:41 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
      ]
     },
     {
@@ -252,52 +258,322 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.688586950302124, validation loss: 0.6886438131332397\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:10| train loss: 0.6998531222343445, validation loss: 0.6888602375984192\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:20| train loss: 0.6862186193466187, validation loss: 0.689547061920166\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:30| train loss: 0.6913824677467346, validation loss: 0.6873278617858887\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:40| train loss: 0.6604955792427063, validation loss: 0.689150333404541\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:50| train loss: 0.6746629476547241, validation loss: 0.6878216862678528\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:60| train loss: 0.7084729075431824, validation loss: 0.6860194802284241\n",
-      "2023-10-10 16:55:17 | INFO     | tempor.models.ts_model:_train:379 | Epoch:70| train loss: 0.6851019263267517, validation loss: 0.683639645576477\n",
-      "2023-10-10 16:55:17 | INFO     | tempor.models.ts_model:_train:379 | Epoch:80| train loss: 0.6695638298988342, validation loss: 0.6823517084121704\n",
-      "2023-10-10 16:55:17 | INFO     | tempor.models.ts_model:_train:379 | Epoch:90| train loss: 0.6456660032272339, validation loss: 0.6871115565299988\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.688586950302124, validation loss: 0.6886438131332397\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:10| train loss: 0.6998531222343445, validation loss: 0.6888602375984192\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:20| train loss: 0.6862186789512634, validation loss: 0.689547061920166\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:30| train loss: 0.6913825273513794, validation loss: 0.6873278617858887\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:40| train loss: 0.6604955196380615, validation loss: 0.6891502141952515\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:50| train loss: 0.6746630072593689, validation loss: 0.6878217458724976\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:60| train loss: 0.7084730267524719, validation loss: 0.6860195398330688\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:70| train loss: 0.6851019263267517, validation loss: 0.6836397647857666\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:80| train loss: 0.6695641279220581, validation loss: 0.6823516488075256\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:90| train loss: 0.6456663608551025, validation loss: 0.6871115565299988\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
      ]
     },
     {
diff --git a/docs/user_guide/usage/tutorial08_benchmarks.ipynb b/docs/user_guide/usage/tutorial08_benchmarks.ipynb
index 68ae1322..1832f6ae 100644
--- a/docs/user_guide/usage/tutorial08_benchmarks.ipynb
+++ b/docs/user_guide/usage/tutorial08_benchmarks.ipynb
@@ -43,14 +43,56 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:56:34 | INFO     | tempor.benchmarks.benchmark:benchmark_models:100 | Test case: model_1\n",
-      "2023-10-10 16:56:36 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.68637615442276, validation loss: 0.6680578589439392\n",
-      "2023-10-10 16:56:36 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6887814402580261, validation loss: 0.6939960718154907\n",
-      "2023-10-10 16:56:37 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6906945109367371, validation loss: 0.6966598033905029\n",
-      "2023-10-10 16:56:37 | INFO     | tempor.benchmarks.benchmark:benchmark_models:100 | Test case: model_2\n",
-      "2023-10-10 16:56:41 | INFO     | tempor.models.ts_ode:_train:607 | Epoch:99| train loss: 0.6465951204299927, validation loss: 0.5632617473602295\n",
-      "2023-10-10 16:56:46 | INFO     | tempor.models.ts_ode:_train:607 | Epoch:99| train loss: 0.913261890411377, validation loss: 0.8132617473602295\n",
-      "2023-10-10 16:56:50 | INFO     | tempor.models.ts_ode:_train:607 | Epoch:99| train loss: 0.7132617831230164, validation loss: 0.8132617473602295\n"
+      "2023-11-16 22:18:41 | INFO     | tempor.benchmarks.benchmark:benchmark_models:104 | Test case: model_1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.68637615442276, validation loss: 0.6680578589439392\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6887814402580261, validation loss: 0.6939960718154907\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6906945109367371, validation loss: 0.6966598033905029\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | tempor.benchmarks.benchmark:benchmark_models:104 | Test case: model_2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:52 | INFO     | tempor.models.ts_ode:_train:617 | Epoch:99| train loss: 0.6465951204299927, validation loss: 0.5632617473602295\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:58 | INFO     | tempor.models.ts_ode:_train:617 | Epoch:99| train loss: 0.913261890411377, validation loss: 0.8132617473602295\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:04 | INFO     | tempor.models.ts_ode:_train:617 | Epoch:99| train loss: 0.7132617831230164, validation loss: 0.8132617473602295\n"
      ]
     },
     {
@@ -87,16 +129,6 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>aucroc</th>\n",
-       "      <td>0.311 +/- 0.083</td>\n",
-       "      <td>0.583 +/- 0.312</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>aucprc</th>\n",
-       "      <td>0.385 +/- 0.028</td>\n",
-       "      <td>0.588 +/- 0.292</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
        "      <th>accuracy</th>\n",
        "      <td>0.602 +/- 0.033</td>\n",
        "      <td>0.519 +/- 0.105</td>\n",
@@ -127,61 +159,71 @@
        "      <td>0.0 +/- 0.0</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>precision_micro</th>\n",
+       "      <th>recall_micro</th>\n",
        "      <td>0.602 +/- 0.033</td>\n",
        "      <td>0.519 +/- 0.105</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>precision_macro</th>\n",
-       "      <td>0.301 +/- 0.016</td>\n",
-       "      <td>0.259 +/- 0.053</td>\n",
+       "      <th>recall_macro</th>\n",
+       "      <td>0.5 +/- 0.0</td>\n",
+       "      <td>0.5 +/- 0.0</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>precision_weighted</th>\n",
-       "      <td>0.363 +/- 0.039</td>\n",
-       "      <td>0.28 +/- 0.104</td>\n",
+       "      <th>recall_weighted</th>\n",
+       "      <td>0.602 +/- 0.033</td>\n",
+       "      <td>0.519 +/- 0.105</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_micro</th>\n",
+       "      <th>precision_micro</th>\n",
        "      <td>0.602 +/- 0.033</td>\n",
        "      <td>0.519 +/- 0.105</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_macro</th>\n",
-       "      <td>0.5 +/- 0.0</td>\n",
-       "      <td>0.5 +/- 0.0</td>\n",
+       "      <th>precision_macro</th>\n",
+       "      <td>0.301 +/- 0.016</td>\n",
+       "      <td>0.259 +/- 0.053</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_weighted</th>\n",
-       "      <td>0.602 +/- 0.033</td>\n",
-       "      <td>0.519 +/- 0.105</td>\n",
+       "      <th>precision_weighted</th>\n",
+       "      <td>0.363 +/- 0.039</td>\n",
+       "      <td>0.28 +/- 0.104</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>mcc</th>\n",
        "      <td>0.0 +/- 0.0</td>\n",
        "      <td>0.0 +/- 0.0</td>\n",
        "    </tr>\n",
+       "    <tr>\n",
+       "      <th>aucprc</th>\n",
+       "      <td>0.385 +/- 0.028</td>\n",
+       "      <td>0.588 +/- 0.292</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>aucroc</th>\n",
+       "      <td>0.311 +/- 0.083</td>\n",
+       "      <td>0.583 +/- 0.312</td>\n",
+       "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "                            model_1          model_2\n",
-       "aucroc              0.311 +/- 0.083  0.583 +/- 0.312\n",
-       "aucprc              0.385 +/- 0.028  0.588 +/- 0.292\n",
        "accuracy            0.602 +/- 0.033  0.519 +/- 0.105\n",
        "f1_score_micro      0.602 +/- 0.033  0.519 +/- 0.105\n",
        "f1_score_macro      0.375 +/- 0.013  0.338 +/- 0.048\n",
        "f1_score_weighted    0.453 +/- 0.04  0.361 +/- 0.116\n",
        "kappa                   0.0 +/- 0.0      0.0 +/- 0.0\n",
        "kappa_quadratic         0.0 +/- 0.0      0.0 +/- 0.0\n",
-       "precision_micro     0.602 +/- 0.033  0.519 +/- 0.105\n",
-       "precision_macro     0.301 +/- 0.016  0.259 +/- 0.053\n",
-       "precision_weighted  0.363 +/- 0.039   0.28 +/- 0.104\n",
        "recall_micro        0.602 +/- 0.033  0.519 +/- 0.105\n",
        "recall_macro            0.5 +/- 0.0      0.5 +/- 0.0\n",
        "recall_weighted     0.602 +/- 0.033  0.519 +/- 0.105\n",
-       "mcc                     0.0 +/- 0.0      0.0 +/- 0.0"
+       "precision_micro     0.602 +/- 0.033  0.519 +/- 0.105\n",
+       "precision_macro     0.301 +/- 0.016  0.259 +/- 0.053\n",
+       "precision_weighted  0.363 +/- 0.039   0.28 +/- 0.104\n",
+       "mcc                     0.0 +/- 0.0      0.0 +/- 0.0\n",
+       "aucprc              0.385 +/- 0.028  0.588 +/- 0.292\n",
+       "aucroc              0.311 +/- 0.083  0.583 +/- 0.312"
       ]
      },
      "metadata": {},
@@ -223,16 +265,6 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>aucroc</th>\n",
-       "      <td>0.311111</td>\n",
-       "      <td>0.083148</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>aucprc</th>\n",
-       "      <td>0.385185</td>\n",
-       "      <td>0.028154</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
        "      <th>accuracy</th>\n",
        "      <td>0.601852</td>\n",
        "      <td>0.032736</td>\n",
@@ -263,6 +295,21 @@
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
+       "      <th>recall_micro</th>\n",
+       "      <td>0.601852</td>\n",
+       "      <td>0.032736</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>recall_macro</th>\n",
+       "      <td>0.500000</td>\n",
+       "      <td>0.000000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>recall_weighted</th>\n",
+       "      <td>0.601852</td>\n",
+       "      <td>0.032736</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
        "      <th>precision_micro</th>\n",
        "      <td>0.601852</td>\n",
        "      <td>0.032736</td>\n",
@@ -278,24 +325,19 @@
        "      <td>0.038647</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_micro</th>\n",
-       "      <td>0.601852</td>\n",
-       "      <td>0.032736</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>recall_macro</th>\n",
-       "      <td>0.500000</td>\n",
+       "      <th>mcc</th>\n",
+       "      <td>0.000000</td>\n",
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_weighted</th>\n",
-       "      <td>0.601852</td>\n",
-       "      <td>0.032736</td>\n",
+       "      <th>aucprc</th>\n",
+       "      <td>0.385185</td>\n",
+       "      <td>0.028154</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>mcc</th>\n",
-       "      <td>0.000000</td>\n",
-       "      <td>0.000000</td>\n",
+       "      <th>aucroc</th>\n",
+       "      <td>0.311111</td>\n",
+       "      <td>0.083148</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -303,21 +345,21 @@
       ],
       "text/plain": [
        "                        mean    stddev\n",
-       "aucroc              0.311111  0.083148\n",
-       "aucprc              0.385185  0.028154\n",
        "accuracy            0.601852  0.032736\n",
        "f1_score_micro      0.601852  0.032736\n",
        "f1_score_macro      0.375458  0.012951\n",
        "f1_score_weighted   0.452788  0.039572\n",
        "kappa               0.000000  0.000000\n",
        "kappa_quadratic     0.000000  0.000000\n",
-       "precision_micro     0.601852  0.032736\n",
-       "precision_macro     0.300926  0.016368\n",
-       "precision_weighted  0.363297  0.038647\n",
        "recall_micro        0.601852  0.032736\n",
        "recall_macro        0.500000  0.000000\n",
        "recall_weighted     0.601852  0.032736\n",
-       "mcc                 0.000000  0.000000"
+       "precision_micro     0.601852  0.032736\n",
+       "precision_macro     0.300926  0.016368\n",
+       "precision_weighted  0.363297  0.038647\n",
+       "mcc                 0.000000  0.000000\n",
+       "aucprc              0.385185  0.028154\n",
+       "aucroc              0.311111  0.083148"
       ]
      },
      "metadata": {},
@@ -357,16 +399,6 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>aucroc</th>\n",
-       "      <td>0.583333</td>\n",
-       "      <td>0.311805</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>aucprc</th>\n",
-       "      <td>0.587731</td>\n",
-       "      <td>0.291568</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
        "      <th>accuracy</th>\n",
        "      <td>0.518519</td>\n",
        "      <td>0.105369</td>\n",
@@ -397,6 +429,21 @@
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
+       "      <th>recall_micro</th>\n",
+       "      <td>0.518519</td>\n",
+       "      <td>0.105369</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>recall_macro</th>\n",
+       "      <td>0.500000</td>\n",
+       "      <td>0.000000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>recall_weighted</th>\n",
+       "      <td>0.518519</td>\n",
+       "      <td>0.105369</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
        "      <th>precision_micro</th>\n",
        "      <td>0.518519</td>\n",
        "      <td>0.105369</td>\n",
@@ -412,24 +459,19 @@
        "      <td>0.104057</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_micro</th>\n",
-       "      <td>0.518519</td>\n",
-       "      <td>0.105369</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>recall_macro</th>\n",
-       "      <td>0.500000</td>\n",
+       "      <th>mcc</th>\n",
+       "      <td>0.000000</td>\n",
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_weighted</th>\n",
-       "      <td>0.518519</td>\n",
-       "      <td>0.105369</td>\n",
+       "      <th>aucprc</th>\n",
+       "      <td>0.587731</td>\n",
+       "      <td>0.291568</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>mcc</th>\n",
-       "      <td>0.000000</td>\n",
-       "      <td>0.000000</td>\n",
+       "      <th>aucroc</th>\n",
+       "      <td>0.583333</td>\n",
+       "      <td>0.311805</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -437,21 +479,21 @@
       ],
       "text/plain": [
        "                        mean    stddev\n",
-       "aucroc              0.583333  0.311805\n",
-       "aucprc              0.587731  0.291568\n",
        "accuracy            0.518519  0.105369\n",
        "f1_score_micro      0.518519  0.105369\n",
        "f1_score_macro      0.338162  0.047609\n",
        "f1_score_weighted   0.360713  0.115623\n",
        "kappa               0.000000  0.000000\n",
        "kappa_quadratic     0.000000  0.000000\n",
-       "precision_micro     0.518519  0.105369\n",
-       "precision_macro     0.259259  0.052684\n",
-       "precision_weighted  0.279964  0.104057\n",
        "recall_micro        0.518519  0.105369\n",
        "recall_macro        0.500000  0.000000\n",
        "recall_weighted     0.518519  0.105369\n",
-       "mcc                 0.000000  0.000000"
+       "precision_micro     0.518519  0.105369\n",
+       "precision_macro     0.259259  0.052684\n",
+       "precision_weighted  0.279964  0.104057\n",
+       "mcc                 0.000000  0.000000\n",
+       "aucprc              0.587731  0.291568\n",
+       "aucroc              0.583333  0.311805"
       ]
      },
      "metadata": {},
diff --git a/docs/user_guide/usage/tutorial09_automl.ipynb b/docs/user_guide/usage/tutorial09_automl.ipynb
index 05b0ea92..e021cb03 100644
--- a/docs/user_guide/usage/tutorial09_automl.ipynb
+++ b/docs/user_guide/usage/tutorial09_automl.ipynb
@@ -61,10 +61,28 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator cde_classifier.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator ode_classifier.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator nn_classifier.\n"
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator cde_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator ode_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator nn_classifier.\n"
      ]
     }
    ],
@@ -120,10 +138,28 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator cde_classifier.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator ode_classifier.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator nn_classifier.\n"
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator cde_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator ode_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator nn_classifier.\n"
      ]
     }
    ],
@@ -178,82 +214,310 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator 'cde_classifier' 1/3.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator 'cde_classifier' 1/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from CDEClassifier:\n",
-      "{'n_iter': 2, 'n_temporal_units_hidden': 13, 'lr': 0.01}\n",
-      "2023-10-10 16:57:31 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_temporal_units_hidden': 13, 'lr': 0.01}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:48 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from CDEClassifier:\n",
-      "{'n_iter': 2, 'n_temporal_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:34 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_temporal_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:53 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from CDEClassifier:\n",
-      "{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}\n",
-      "2023-10-10 16:57:38 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator 'ode_classifier' 2/3.\n",
-      "2023-10-10 16:57:38 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:57:38 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:58 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator 'ode_classifier' 2/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:58 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:58 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from ODEClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 13, 'lr': 0.01}\n",
-      "2023-10-10 16:57:41 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 13, 'lr': 0.01}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:01 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from ODEClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:44 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:05 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from ODEClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 20, 'lr': 0.001}\n",
-      "2023-10-10 16:57:48 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator 'nn_classifier' 3/3.\n",
-      "2023-10-10 16:57:48 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:57:48 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 20, 'lr': 0.001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:08 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator 'nn_classifier' 3/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:08 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:08 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from NeuralNetClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 13, 'lr': 0.01}\n",
-      "2023-10-10 16:57:48 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.6882655620574951\n",
-      "2023-10-10 16:57:49 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6865323185920715\n",
-      "2023-10-10 16:57:49 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6861070990562439\n",
-      "2023-10-10 16:57:50 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6825112104415894\n",
-      "2023-10-10 16:57:50 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6819757223129272\n",
-      "2023-10-10 16:57:50 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 13, 'lr': 0.01}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:09 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.6882655620574951\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:09 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6865323185920715\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:10 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6861070990562439\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:10 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6825112104415894\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:11 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6819757223129272\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:11 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from NeuralNetClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:51 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.691270649433136\n",
-      "2023-10-10 16:57:52 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6910805106163025\n",
-      "2023-10-10 16:57:52 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6900344491004944\n",
-      "2023-10-10 16:57:53 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6906693577766418\n",
-      "2023-10-10 16:57:53 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6910874843597412\n",
-      "2023-10-10 16:57:53 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:12 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.691270649433136\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:12 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6910805106163025\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:13 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6900344491004944\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:13 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6906693577766418\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:14 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6910874843597412\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:14 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from NeuralNetClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 20, 'lr': 0.001}\n",
-      "2023-10-10 16:57:54 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.6896671056747437\n",
-      "2023-10-10 16:57:54 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6886401772499084\n",
-      "2023-10-10 16:57:55 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6884845495223999\n",
-      "2023-10-10 16:57:55 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6887120604515076\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6893669962882996\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 20, 'lr': 0.001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:14 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.6896671056747437\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:15 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6886401772499084\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:15 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6884845495223999\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6887120604515076\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6893669962882996\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for cde_classifier scores:\n",
-      "[0.47310606060606064, 0.48396464646464643, 0.4840277777777778].\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "[0.46300505050505053, 0.48396464646464643, 0.4901515151515151].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for ode_classifier scores:\n",
-      "[0.47992424242424236, 0.48409090909090907, 0.484090909090909].\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "[0.47992424242424236, 0.48409090909090907, 0.484090909090909].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for nn_classifier scores:\n",
-      "[0.37064393939393936, 0.5370580808080809, 0.4579545454545454].\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:460 | \n",
+      "[0.37064393939393936, 0.5370580808080809, 0.4579545454545454].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:436 | \n",
       "All estimator definitions searched:\n",
-      "['cde_classifier', 'ode_classifier', 'nn_classifier']\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:461 | \n",
+      "['cde_classifier', 'ode_classifier', 'nn_classifier']\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:437 | \n",
       "Best scores for each estimator searched:\n",
-      "[0.4840277777777778, 0.48409090909090907, 0.5370580808080809]\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:462 | \n",
+      "[0.4901515151515151, 0.48409090909090907, 0.5370580808080809]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:438 | \n",
       "Best hyperparameters for each estimator searched:\n",
-      "[{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}, {'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}, {'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}]\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:587 | \n",
+      "[{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}, {'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}, {'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:565 | \n",
       "Selected score 0.5370580808080809 for nn_classifier with hyperparameters:\n",
-      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:587 | \n",
-      "Selected score 0.48409090909090907 for ode_classifier with hyperparameters:\n",
-      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:587 | \n",
-      "Selected score 0.4840277777777778 for cde_classifier with hyperparameters:\n",
+      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:565 | \n",
+      "Selected score 0.4901515151515151 for cde_classifier with hyperparameters:\n",
       "{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}\n"
      ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:565 | \n",
+      "Selected score 0.48409090909090907 for ode_classifier with hyperparameters:\n",
+      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
     }
    ],
    "source": [
@@ -333,19 +597,19 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ODEClassifier</span><span style=\"font-weight: bold\">(</span>\n",
-       "    <span style=\"color: #808000; text-decoration-color: #808000\">name</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'ode_classifier'</span>,\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">CDEClassifier</span><span style=\"font-weight: bold\">(</span>\n",
+       "    <span style=\"color: #808000; text-decoration-color: #808000\">name</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'cde_classifier'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">category</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'prediction.one_off.classification'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">plugin_type</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'method'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">params</span>=<span style=\"font-weight: bold\">{</span>\n",
-       "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_units_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_units_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">100</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_layers_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'nonlin'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'relu'</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'dropout'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'atol'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'rtol'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'interpolation'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'cubic'</span>,\n",
-       "        <span style=\"color: #008000; text-decoration-color: #008000\">'lr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0001</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'lr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.001</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'weight_decay'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.001</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_iter'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'batch_size'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">500</span>,\n",
@@ -361,19 +625,19 @@
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[1;35mODEClassifier\u001b[0m\u001b[1m(\u001b[0m\n",
-       "    \u001b[33mname\u001b[0m=\u001b[32m'ode_classifier'\u001b[0m,\n",
+       "\u001b[1;35mCDEClassifier\u001b[0m\u001b[1m(\u001b[0m\n",
+       "    \u001b[33mname\u001b[0m=\u001b[32m'cde_classifier'\u001b[0m,\n",
        "    \u001b[33mcategory\u001b[0m=\u001b[32m'prediction.one_off.classification'\u001b[0m,\n",
        "    \u001b[33mplugin_type\u001b[0m=\u001b[32m'method'\u001b[0m,\n",
        "    \u001b[33mparams\u001b[0m=\u001b[1m{\u001b[0m\n",
-       "        \u001b[32m'n_units_hidden'\u001b[0m: \u001b[1;36m11\u001b[0m,\n",
+       "        \u001b[32m'n_units_hidden'\u001b[0m: \u001b[1;36m100\u001b[0m,\n",
        "        \u001b[32m'n_layers_hidden'\u001b[0m: \u001b[1;36m1\u001b[0m,\n",
        "        \u001b[32m'nonlin'\u001b[0m: \u001b[32m'relu'\u001b[0m,\n",
        "        \u001b[32m'dropout'\u001b[0m: \u001b[1;36m0\u001b[0m,\n",
        "        \u001b[32m'atol'\u001b[0m: \u001b[1;36m0.01\u001b[0m,\n",
        "        \u001b[32m'rtol'\u001b[0m: \u001b[1;36m0.01\u001b[0m,\n",
        "        \u001b[32m'interpolation'\u001b[0m: \u001b[32m'cubic'\u001b[0m,\n",
-       "        \u001b[32m'lr'\u001b[0m: \u001b[1;36m0.0001\u001b[0m,\n",
+       "        \u001b[32m'lr'\u001b[0m: \u001b[1;36m0.001\u001b[0m,\n",
        "        \u001b[32m'weight_decay'\u001b[0m: \u001b[1;36m0.001\u001b[0m,\n",
        "        \u001b[32m'n_iter'\u001b[0m: \u001b[1;36m2\u001b[0m,\n",
        "        \u001b[32m'batch_size'\u001b[0m: \u001b[1;36m500\u001b[0m,\n",
@@ -394,19 +658,19 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">CDEClassifier</span><span style=\"font-weight: bold\">(</span>\n",
-       "    <span style=\"color: #808000; text-decoration-color: #808000\">name</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'cde_classifier'</span>,\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ODEClassifier</span><span style=\"font-weight: bold\">(</span>\n",
+       "    <span style=\"color: #808000; text-decoration-color: #808000\">name</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'ode_classifier'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">category</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'prediction.one_off.classification'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">plugin_type</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'method'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">params</span>=<span style=\"font-weight: bold\">{</span>\n",
-       "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_units_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">100</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_units_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_layers_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'nonlin'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'relu'</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'dropout'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'atol'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'rtol'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'interpolation'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'cubic'</span>,\n",
-       "        <span style=\"color: #008000; text-decoration-color: #008000\">'lr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.001</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'lr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0001</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'weight_decay'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.001</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_iter'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'batch_size'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">500</span>,\n",
@@ -422,19 +686,19 @@
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[1;35mCDEClassifier\u001b[0m\u001b[1m(\u001b[0m\n",
-       "    \u001b[33mname\u001b[0m=\u001b[32m'cde_classifier'\u001b[0m,\n",
+       "\u001b[1;35mODEClassifier\u001b[0m\u001b[1m(\u001b[0m\n",
+       "    \u001b[33mname\u001b[0m=\u001b[32m'ode_classifier'\u001b[0m,\n",
        "    \u001b[33mcategory\u001b[0m=\u001b[32m'prediction.one_off.classification'\u001b[0m,\n",
        "    \u001b[33mplugin_type\u001b[0m=\u001b[32m'method'\u001b[0m,\n",
        "    \u001b[33mparams\u001b[0m=\u001b[1m{\u001b[0m\n",
-       "        \u001b[32m'n_units_hidden'\u001b[0m: \u001b[1;36m100\u001b[0m,\n",
+       "        \u001b[32m'n_units_hidden'\u001b[0m: \u001b[1;36m11\u001b[0m,\n",
        "        \u001b[32m'n_layers_hidden'\u001b[0m: \u001b[1;36m1\u001b[0m,\n",
        "        \u001b[32m'nonlin'\u001b[0m: \u001b[32m'relu'\u001b[0m,\n",
        "        \u001b[32m'dropout'\u001b[0m: \u001b[1;36m0\u001b[0m,\n",
        "        \u001b[32m'atol'\u001b[0m: \u001b[1;36m0.01\u001b[0m,\n",
        "        \u001b[32m'rtol'\u001b[0m: \u001b[1;36m0.01\u001b[0m,\n",
        "        \u001b[32m'interpolation'\u001b[0m: \u001b[32m'cubic'\u001b[0m,\n",
-       "        \u001b[32m'lr'\u001b[0m: \u001b[1;36m0.001\u001b[0m,\n",
+       "        \u001b[32m'lr'\u001b[0m: \u001b[1;36m0.0001\u001b[0m,\n",
        "        \u001b[32m'weight_decay'\u001b[0m: \u001b[1;36m0.001\u001b[0m,\n",
        "        \u001b[32m'n_iter'\u001b[0m: \u001b[1;36m2\u001b[0m,\n",
        "        \u001b[32m'batch_size'\u001b[0m: \u001b[1;36m500\u001b[0m,\n",
@@ -493,10 +757,28 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with cde_classifier>.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with ode_classifier>.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with nn_classifier>.\n"
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with cde_classifier>.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with ode_classifier>.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with nn_classifier>.\n"
      ]
     }
    ],
@@ -580,10 +862,28 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with cde_classifier>.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with ode_classifier>.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with nn_classifier>.\n"
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with cde_classifier>.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with ode_classifier>.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with nn_classifier>.\n"
      ]
     }
    ],
@@ -645,80 +945,308 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator '<Pipeline with cde_classifier>' 1/3.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator '<Pipeline with cde_classifier>' 1/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.cde_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 17, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:14 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 17, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:32 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_minmax_scaler->prediction.one_off.classification.cde_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 14, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:20 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 14, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:38 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_standard_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.cde_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 6, 'lr': 0.0001}}}\n",
-      "2023-10-10 16:58:26 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator '<Pipeline with ode_classifier>' 2/3.\n",
-      "2023-10-10 16:58:26 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:58:26 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 6, 'lr': 0.0001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:45 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator '<Pipeline with ode_classifier>' 2/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:45 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:45 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.ode_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 17, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:45 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 17, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:00 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_minmax_scaler->prediction.one_off.classification.ode_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 14, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:50 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 14, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:06 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_standard_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.ode_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 6, 'lr': 0.0001}}}\n",
-      "2023-10-10 16:58:56 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator '<Pipeline with nn_classifier>' 3/3.\n",
-      "2023-10-10 16:58:56 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:58:56 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 6, 'lr': 0.0001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:12 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator '<Pipeline with nn_classifier>' 3/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:12 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:12 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.nn_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 17, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:58 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6918608546257019, validation loss: 0.6910383105278015\n",
-      "2023-10-10 16:59:01 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6909429430961609, validation loss: 0.6872987151145935\n",
-      "2023-10-10 16:59:04 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.692150354385376, validation loss: 0.6879752278327942\n",
-      "2023-10-10 16:59:07 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6926363706588745, validation loss: 0.6881052851676941\n",
-      "2023-10-10 16:59:10 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6915659308433533, validation loss: 0.6888220310211182\n",
-      "2023-10-10 16:59:11 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 17, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:14 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6918608546257019, validation loss: 0.6910383105278015\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:17 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6909429430961609, validation loss: 0.6872987151145935\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:20 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.692150354385376, validation loss: 0.6879752278327942\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:23 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6926363706588745, validation loss: 0.6881052851676941\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:26 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6915659308433533, validation loss: 0.6888220310211182\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:27 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_minmax_scaler->prediction.one_off.classification.nn_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 14, 'lr': 0.001}}}\n",
-      "2023-10-10 16:59:12 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6920319199562073, validation loss: 0.6892786622047424\n",
-      "2023-10-10 16:59:13 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6913070678710938, validation loss: 0.6893746852874756\n",
-      "2023-10-10 16:59:14 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6924894452095032, validation loss: 0.6886722445487976\n",
-      "2023-10-10 16:59:15 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6930248141288757, validation loss: 0.6894596815109253\n",
-      "2023-10-10 16:59:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6921922564506531, validation loss: 0.6893032789230347\n",
-      "2023-10-10 16:59:16 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 14, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:27 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6920319199562073, validation loss: 0.6892786622047424\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:28 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6913070678710938, validation loss: 0.6893746852874756\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:30 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6924894452095032, validation loss: 0.6886722445487976\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:31 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6930248141288757, validation loss: 0.6894596815109253\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:32 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6921922564506531, validation loss: 0.6893032789230347\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:32 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_standard_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.nn_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 6, 'lr': 0.0001}}}\n",
-      "2023-10-10 16:59:17 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6914154291152954, validation loss: 0.6911771297454834\n",
-      "2023-10-10 16:59:18 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.690178394317627, validation loss: 0.6934390664100647\n",
-      "2023-10-10 16:59:19 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6933540105819702, validation loss: 0.6896734833717346\n",
-      "2023-10-10 16:59:20 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.694333553314209, validation loss: 0.690070390701294\n",
-      "2023-10-10 16:59:21 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6921719312667847, validation loss: 0.6909489631652832\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 6, 'lr': 0.0001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:33 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6914154291152954, validation loss: 0.6911771297454834\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:34 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.690178394317627, validation loss: 0.6934390664100647\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:35 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6933540105819702, validation loss: 0.6896734833717346\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:36 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.694333553314209, validation loss: 0.690070390701294\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6921719312667847, validation loss: 0.6909489631652832\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for <Pipeline with cde_classifier> scores:\n",
-      "[0.4716540404040404, 0.47575757575757577, 0.4657828282828283].\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "[0.4736742424242424, 0.4737373737373737, 0.4657828282828283].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for <Pipeline with ode_classifier> scores:\n",
-      "[0.47575757575757577, 0.4779040404040404, 0.4779671717171716].\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "[0.47575757575757577, 0.4779040404040404, 0.4779671717171716].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for <Pipeline with nn_classifier> scores:\n",
-      "[0.48478535353535346, 0.4361742424242424, 0.5071338383838383].\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:460 | \n",
+      "[0.48478535353535346, 0.4361742424242424, 0.5071338383838383].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:436 | \n",
       "All estimator definitions searched:\n",
-      "['<Pipeline with cde_classifier>', '<Pipeline with ode_classifier>', '<Pipeline with nn_classifier>']\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:461 | \n",
+      "['<Pipeline with cde_classifier>', '<Pipeline with ode_classifier>', '<Pipeline with nn_classifier>']\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:437 | \n",
       "Best scores for each estimator searched:\n",
-      "[0.47575757575757577, 0.4779671717171716, 0.5071338383838383]\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:462 | \n",
+      "[0.4737373737373737, 0.4779671717171716, 0.5071338383838383]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:438 | \n",
       "Best hyperparameters for each estimator searched:\n",
-      "[{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'mean', '<candidates>(preprocessing.scaling.static)': 'static_minmax_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_minmax_scaler', '[ts_minmax_scaler](clip)': False, '[cde_classifier](n_iter)': 2, '[cde_classifier](n_temporal_units_hidden)': 14, '[cde_classifier](lr)': 0.001}, {'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[ode_classifier](n_iter)': 2, '[ode_classifier](n_units_hidden)': 6, '[ode_classifier](lr)': 0.0001}, {'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[nn_classifier](n_iter)': 2, '[nn_classifier](n_units_hidden)': 6, '[nn_classifier](lr)': 0.0001}]\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:768 | \n",
+      "[{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'mean', '<candidates>(preprocessing.scaling.static)': 'static_minmax_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_minmax_scaler', '[ts_minmax_scaler](clip)': False, '[cde_classifier](n_iter)': 2, '[cde_classifier](n_temporal_units_hidden)': 14, '[cde_classifier](lr)': 0.001}, {'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[ode_classifier](n_iter)': 2, '[ode_classifier](n_units_hidden)': 6, '[ode_classifier](lr)': 0.0001}, {'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[nn_classifier](n_iter)': 2, '[nn_classifier](n_units_hidden)': 6, '[nn_classifier](lr)': 0.0001}]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:748 | \n",
       "Selected score 0.5071338383838383 for <Pipeline with nn_classifier> with hyperparameters:\n",
-      "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[nn_classifier](n_iter)': 2, '[nn_classifier](n_units_hidden)': 6, '[nn_classifier](lr)': 0.0001}\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:768 | \n",
+      "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[nn_classifier](n_iter)': 2, '[nn_classifier](n_units_hidden)': 6, '[nn_classifier](lr)': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:748 | \n",
       "Selected score 0.4779671717171716 for <Pipeline with ode_classifier> with hyperparameters:\n",
-      "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[ode_classifier](n_iter)': 2, '[ode_classifier](n_units_hidden)': 6, '[ode_classifier](lr)': 0.0001}\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:768 | \n",
-      "Selected score 0.47575757575757577 for <Pipeline with cde_classifier> with hyperparameters:\n",
+      "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[ode_classifier](n_iter)': 2, '[ode_classifier](n_units_hidden)': 6, '[ode_classifier](lr)': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:748 | \n",
+      "Selected score 0.4737373737373737 for <Pipeline with cde_classifier> with hyperparameters:\n",
       "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'mean', '<candidates>(preprocessing.scaling.static)': 'static_minmax_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_minmax_scaler', '[ts_minmax_scaler](clip)': False, '[cde_classifier](n_iter)': 2, '[cde_classifier](n_temporal_units_hidden)': 14, '[cde_classifier](lr)': 0.001}\n"
      ]
     }
@@ -986,10 +1514,28 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator cde_classifier.\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator ode_classifier.\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator nn_classifier.\n"
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator cde_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator ode_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator nn_classifier.\n"
      ]
     }
    ],
diff --git a/pypi.md b/pypi.md
index ad66a251..bbfa6843 100644
--- a/pypi.md
+++ b/pypi.md
@@ -472,6 +472,7 @@ Prediction where targets are temporal (time series).
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial01_custom_method.ipynb) - [Writing a Custom Method Plugin](https://github.com/vanderschaarlab/temporai/tree/main/tutorials/extending/tutorial01_custom_method.ipynb)
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial02_testing_custom_method.ipynb) - [Testing a Custom Method Plugin](https://github.com/vanderschaarlab/temporai/tree/main/tutorials/extending/tutorial02_testing_custom_method.ipynb)
 - [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial03_custom_datasource.ipynb) - [Writing a Custom Data Source Plugin](https://github.com/vanderschaarlab/temporai/tree/main/tutorials/extending/tutorial03_custom_datasource.ipynb)
+- [![Test In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/vanderschaarlab/temporai/blob/main/tutorials/extending/tutorial04_custom_metric.ipynb) - [Writing a Custom Metric Plugin](https://github.com/vanderschaarlab/temporai/tree/main/tutorials/extending/tutorial04_custom_metric.ipynb)
 
 
 
diff --git a/src/tempor/__init__.py b/src/tempor/__init__.py
index f09bb975..38c36502 100644
--- a/src/tempor/__init__.py
+++ b/src/tempor/__init__.py
@@ -34,6 +34,7 @@
 
 import tempor.datasources  # noqa: E402 F401
 import tempor.methods  # noqa: E402 F401
+import tempor.metrics  # noqa: E402 F401
 
 # Prepare the plugin loader:
 from tempor.core import plugins  # noqa: E402
diff --git a/src/tempor/automl/__init__.py b/src/tempor/automl/__init__.py
index 39fb4841..20c41d93 100644
--- a/src/tempor/automl/__init__.py
+++ b/src/tempor/automl/__init__.py
@@ -1,7 +1,1 @@
 """Package directory with the AutoML functionality."""
-
-from ._types import OptimDirection
-
-__all__ = [
-    "OptimDirection",
-]
diff --git a/src/tempor/automl/_types.py b/src/tempor/automl/_types.py
index 02285471..2a8ac4c5 100644
--- a/src/tempor/automl/_types.py
+++ b/src/tempor/automl/_types.py
@@ -2,15 +2,10 @@
 
 from typing import Type, Union
 
-from typing_extensions import Literal
-
 from tempor.methods.core import BasePredictor
 
 from .pipeline_selector import PipelineSelector
 
-OptimDirection = Literal["minimize", "maximize"]
-"""Optimization direction for AutoML (with respect to a metric)"""
-
 AutoMLCompatibleEstimator = Union[Type[BasePredictor], PipelineSelector]
 """Estimator types supported by AutoML. Either a predictor class,
 or a special `PipelineSelector` object for pipeline search.
diff --git a/src/tempor/automl/seeker.py b/src/tempor/automl/seeker.py
index 412a151e..8a46e449 100644
--- a/src/tempor/automl/seeker.py
+++ b/src/tempor/automl/seeker.py
@@ -4,7 +4,7 @@
 import copy
 import functools
 import warnings
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast
 
 import numpy as np
 import optuna
@@ -20,8 +20,9 @@
 from tempor.log import logger
 from tempor.methods.core import BasePredictor
 from tempor.methods.core.params import Params
+from tempor.metrics import metric_typing
 
-from ._types import AutoMLCompatibleEstimator, OptimDirection
+from ._types import AutoMLCompatibleEstimator
 from .pipeline_selector import (
     DEFAULT_STATIC_IMPUTERS,
     DEFAULT_STATIC_SCALERS,
@@ -48,13 +49,6 @@
   - `"grid"`: Use a tuner based on `optuna.samplers.GridSampler`.
 """
 
-SupportedMetric = Union[
-    evaluation.ClassifierSupportedMetric,
-    evaluation.RegressionSupportedMetric,
-    evaluation.TimeToEventSupportedMetric,
-]
-"""The type denoting all metrics supported by ``Seeker`` classes"""
-
 TUNER_OPTUNA_SAMPLER_MAP: Dict[TunerType, Any] = {
     "bayesian": optuna.samplers.TPESampler,
     "random": optuna.samplers.RandomSampler,
@@ -64,33 +58,6 @@
 }
 """A map from `TunerType` to the corresponding `optuna` sampler class"""
 
-METRIC_DIRECTION_MAP: Dict[SupportedMetric, OptimDirection] = {
-    # ClassifierSupportedMetric:
-    "aucroc": "maximize",
-    "aucprc": "maximize",
-    "accuracy": "maximize",
-    "f1_score_micro": "maximize",
-    "f1_score_macro": "maximize",
-    "f1_score_weighted": "maximize",
-    "kappa": "maximize",
-    "kappa_quadratic": "maximize",
-    "precision_micro": "maximize",
-    "precision_macro": "maximize",
-    "precision_weighted": "maximize",
-    "recall_micro": "maximize",
-    "recall_macro": "maximize",
-    "recall_weighted": "maximize",
-    "mcc": "maximize",
-    # RegressionSupportedMetric:
-    "mse": "minimize",
-    "mae": "minimize",
-    "r2": "maximize",
-    # TimeToEventMetricCallable:
-    "c_index": "maximize",
-    "brier_score": "minimize",
-}
-"""A map from metric (`SupportedMetric`) to its optimization direction (`OptimDirection`)"""
-
 # TODO: Parallelism support (per-estimator).
 
 
@@ -98,7 +65,7 @@ def evaluation_callback_dispatch(
     estimator: Type[BasePredictor],
     dataset: PredictiveDataset,
     task_type: PredictiveTaskType,
-    metric: SupportedMetric,
+    metric: str,
     n_cv_folds: int,
     random_state: int,
     horizon: Optional[data_typing.TimeIndex],
@@ -117,7 +84,7 @@ def evaluation_callback_dispatch(
             The dataset to use.
         task_type (PredictiveTaskType):
             The task type of the predictor.
-        metric (SupportedMetric):
+        metric (str):
             The metric to be used for evaluation.
         n_cv_folds (int):
             Number of cross-validation folds to use.
@@ -199,7 +166,7 @@ def __init__(  # pylint: disable=unused-argument
         task_type: PredictiveTaskType,
         estimator_names: List[str],
         estimator_defs: List[Any],
-        metric: SupportedMetric,
+        metric: str,
         dataset: PredictiveDataset,
         *,
         return_top_k: int = 3,
@@ -230,7 +197,7 @@ def __init__(  # pylint: disable=unused-argument
                 Friendly names of estimators. Will be passed one-by-one to ``_init_estimator`` method calls.
             estimator_defs (List[Any]):
                 Definition of estimators. Will be passed one-by-one to ``_init_estimator`` method calls.
-            metric (SupportedMetric):
+            metric (str):
                 The metric to use for evaluation.
             dataset (PredictiveDataset):
                 The dataset to use for evaluation.
@@ -280,7 +247,7 @@ def __init__(  # pylint: disable=unused-argument
 
         self.study_name = study_name
         self.task_type: PredictiveTaskType = task_type
-        self.metric: SupportedMetric = metric
+        self.metric = metric
         self.dataset = dataset
         self.return_top_k = return_top_k
         self.num_cv_folds = num_cv_folds
@@ -295,6 +262,8 @@ def __init__(  # pylint: disable=unused-argument
         self.raise_exceptions = raise_exceptions
         self.silence_warnings = silence_warnings
 
+        self._metric_object = plugin_loader.get(f"{task_type}.{self.metric}", plugin_type="metric")
+
         if len(estimator_defs) != len(estimator_names):
             raise ValueError("`estimator_defs` and `estimator_names` must be the same length.")
         self.estimator_names = estimator_names
@@ -324,7 +293,7 @@ def __init__(  # pylint: disable=unused-argument
                 raise ValueError("Passing a custom tuner with `optuna.samplers.GridSampler` is not supported")
         self.custom_tuner = custom_tuner
 
-        self.direction: OptimDirection = METRIC_DIRECTION_MAP[self.metric]
+        self.direction: metric_typing.MetricDirection = self._metric_object.direction
         self.estimators: List[AutoMLCompatibleEstimator] = []
         self.tuners: List[BaseTuner] = []
 
@@ -493,7 +462,7 @@ def __init__(
         study_name: str,
         task_type: PredictiveTaskType,
         estimator_names: List[str],
-        metric: SupportedMetric,
+        metric: str,
         dataset: PredictiveDataset,
         *,
         return_top_k: int = 3,
@@ -523,7 +492,7 @@ def __init__(
             estimator_names (List[str]):
                 The candidate predictors. Provide plugin names (without category qualification), e.g. like
                 ``["nn_classifier", "cde_classifier"]``.
-            metric (SupportedMetric):
+            metric (str):
                 See `~tempor.automl.seeker.BaseSeeker`.
             dataset (PredictiveDataset):
                 See `~tempor.automl.seeker.BaseSeeker`.
@@ -604,7 +573,7 @@ def __init__(
         study_name: str,
         task_type: PredictiveTaskType,
         estimator_names: List[str],
-        metric: SupportedMetric,
+        metric: str,
         dataset: PredictiveDataset,
         *,
         static_imputers: List[str] = DEFAULT_STATIC_IMPUTERS,
@@ -650,7 +619,7 @@ def __init__(
             estimator_names (List[str]):
                 The candidate predictors that will be the last step of the pipeline. Provide plugin names
                 (without category qualification), e.g. like ``["nn_classifier", "cde_classifier"]``.
-            metric (SupportedMetric):
+            metric (str):
                 See `~tempor.automl.seeker.BaseSeeker`.
             dataset (PredictiveDataset):
                 See `~tempor.automl.seeker.BaseSeeker`.
diff --git a/src/tempor/automl/tuner.py b/src/tempor/automl/tuner.py
index 6b305908..d044b470 100644
--- a/src/tempor/automl/tuner.py
+++ b/src/tempor/automl/tuner.py
@@ -13,8 +13,9 @@
 from tempor.log import logger
 from tempor.methods.core._base_predictor import BasePredictor
 from tempor.methods.core.params import Params
+from tempor.metrics import metric_typing
 
-from ._types import AutoMLCompatibleEstimator, OptimDirection
+from ._types import AutoMLCompatibleEstimator
 from .pipeline_selector import PipelineSelector
 
 # TODO: Handle other hyperparameter tuning frameworks, e.g. hyperband.
@@ -56,7 +57,7 @@ class BaseTuner(abc.ABC):
     def __init__(  # pylint: disable=unused-argument
         self,
         study_name: str,
-        direction: OptimDirection,
+        direction: metric_typing.MetricDirection,
         **kwargs: Any,
     ):
         """Base hyperparameter tuner from which tuner implementations should derive. Defines the initializer and the
@@ -65,7 +66,7 @@ def __init__(  # pylint: disable=unused-argument
         Args:
             study_name (str):
                 Study name.
-            direction (OptimDirection):
+            direction (metric_typing.MetricDirection):
                 Optimization direction (`"minimize"` or `"maximize"`).
             **kwargs (Any):
                 Currently unused.
@@ -115,7 +116,7 @@ class OptunaTuner(BaseTuner):
     def __init__(
         self,
         study_name: str,
-        direction: OptimDirection,
+        direction: metric_typing.MetricDirection,
         *,
         study_sampler: optuna.samplers.BaseSampler,
         study_storage: Optional[Union[str, optuna.storages.BaseStorage]] = None,
@@ -129,7 +130,7 @@ def __init__(
         Args:
             study_name (str):
                 Study name.
-            direction (OptimDirection):
+            direction (metric_typing.MetricDirection):
                 Optimization direction (`"minimize"` or `"maximize"`).
             study_sampler (optuna.samplers.BaseSampler):
                 An `optuna` sampler (passed to `optuna.create_study`).
diff --git a/src/tempor/benchmarks/__init__.py b/src/tempor/benchmarks/__init__.py
index 4d87b0a0..49a6742c 100644
--- a/src/tempor/benchmarks/__init__.py
+++ b/src/tempor/benchmarks/__init__.py
@@ -2,14 +2,12 @@
 
 from .benchmark import benchmark_models, visualize_benchmark  # noqa: F401
 from .evaluation import (  # noqa: F401
-    ClassifierSupportedMetric,
     OutputMetric,
-    RegressionSupportedMetric,
-    classifier_supported_metrics,
+    builtin_metrics_prediction_oneoff_classification,
+    builtin_metrics_prediction_oneoff_regression,
+    builtin_metrics_time_to_event,
     evaluate_prediction_oneoff_classifier,
     evaluate_prediction_oneoff_regressor,
     evaluate_time_to_event,
     output_metrics,
-    regression_supported_metrics,
-    time_to_event_supported_metrics,
 )
diff --git a/src/tempor/benchmarks/benchmark.py b/src/tempor/benchmarks/benchmark.py
index 1d8147e1..26e12984 100644
--- a/src/tempor/benchmarks/benchmark.py
+++ b/src/tempor/benchmarks/benchmark.py
@@ -179,7 +179,7 @@ def visualize_benchmark(results: Dict[str, pd.DataFrame], palette: str = "viridi
                 # Known issue with seaborn 0.13.0+. If the y values are non-unique, it seems to expect the yerr values
                 # only as many as there are unique y values. We hence remap using numpy.unique with index return.
                 values = df_sns[df_sns["metric"] == metric]["mean"]
-                _, unique_indexes = np.unique(values, return_index=True)
+                _, unique_indexes = np.unique(values, return_index=True)  # pyright: ignore
                 err_remapped = err.loc[metric, :][unique_indexes]
                 out = sns.barplot(
                     df_sns[df_sns["metric"] == metric],
diff --git a/src/tempor/benchmarks/evaluation.py b/src/tempor/benchmarks/evaluation.py
index 69241b35..d6b143ab 100644
--- a/src/tempor/benchmarks/evaluation.py
+++ b/src/tempor/benchmarks/evaluation.py
@@ -3,7 +3,7 @@
 import copy
 import warnings
 from time import time
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Sequence, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Tuple, cast
 
 import numpy as np
 import pandas as pd
@@ -13,12 +13,12 @@
 import sklearn.model_selection
 from typing_extensions import Literal, get_args
 
-from tempor.core import pydantic_utils
+from tempor.core import plugins, pydantic_utils
 from tempor.data import data_typing, dataset, samples
 from tempor.log import logger
+from tempor.metrics import metric as metric_module
 from tempor.models.utils import enable_reproducibility
 
-from . import metrics as tempor_metrics
 from . import utils
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -28,100 +28,15 @@
 
 # TODO: Benchmarking workflow for missing cases.
 
-ClassifierSupportedMetric = Literal[
-    "aucroc",
-    "aucprc",
-    "accuracy",
-    "f1_score_micro",
-    "f1_score_macro",
-    "f1_score_weighted",
-    "kappa",
-    "kappa_quadratic",
-    "precision_micro",
-    "precision_macro",
-    "precision_weighted",
-    "recall_micro",
-    "recall_macro",
-    "recall_weighted",
-    "mcc",
+_plugin_loader = plugins.PluginLoader()
+builtin_metrics_prediction_oneoff_classification = _plugin_loader.list(plugin_type="metric")["prediction"]["one_off"][
+    "classification"
 ]
-"""Evaluation metrics supported in the classification task setting.
-
-Possible values:
-    - ``"aucroc"``:
-        The Area Under the Receiver Operating Characteristic Curve (ROC AUC) from prediction scores.
-    - ``"aucprc"``:
-        The average precision summarizes a precision-recall curve as the weighted mean of precisions achieved at each
-        threshold, with the increase in recall from the previous threshold used as the weight.
-    - ``"accuracy"``:
-        Accuracy classification score.
-    - ``"f1_score_micro"``:
-        F1 score is a harmonic mean of the precision and recall. This version uses the ``"micro"`` average:
-        calculate metrics globally by counting the total true positives, false negatives and false positives.
-    - ``"f1_score_macro"``:
-        F1 score is a harmonic mean of the precision and recall. This version uses the ``"macro"`` average:
-        calculate metrics for each label, and find their unweighted mean. This does not take label imbalance into
-        account.
-    - ``"f1_score_weighted"``:
-        F1 score is a harmonic mean of the precision and recall. This version uses the "weighted" average:
-        Calculate metrics for each label, and find their average weighted by support
-        (the number of true instances for each label).
-    - ``"kappa"``, ``"kappa_quadratic"``:
-        Computes Cohen's kappa, a score that expresses the level of agreement between two annotators on a
-        classification problem.
-    - ``"precision_micro"``:
-        Precision is defined as the number of true positives over the number of true positives plus the number of false
-        positives. This version(micro) calculates metrics globally by counting the total true positives.
-    - ``"precision_macro"``:
-        Precision is defined as the number of true positives over the number of true positives plus the number of
-        false positives. This version (macro) calculates metrics for each label, and finds their unweighted mean.
-    - ``"precision_weighted"``:
-        Precision is defined as the number of true positives over the number of true positives plus the number of
-        false positives. This version (weighted) calculates metrics for each label, and find their average weighted
-        by support.
-    - ``"recall_micro"``:
-        Recall is defined as the number of true positives over the number of true positives plus the number of false
-        negatives. This version (micro) calculates metrics globally by counting the total true positives.
-    - ``"recall_macro"``:
-        Recall is defined as the number of true positives over the number of true positives plus the number of false
-        negatives. This version (macro) calculates metrics for each label, and finds their unweighted mean.
-    - ``"recall_weighted"``:
-        Recall is defined as the number of true positives over the number of true positives plus the number of false
-        negatives. This version(weighted) calculates metrics for each label, and find their average weighted by support.
-    - ``"mcc"``:
-        The Matthews Correlation Coefficient is used in machine learning as a measure of the quality of binary and
-        multiclass classifications. It takes into account true and false positives and negatives and is generally
-        regarded as a balanced measure which can be used even if the classes are of very different sizes.
-"""
-
-RegressionSupportedMetric = Literal[
-    "mse",
-    "mae",
-    "r2",
+builtin_metrics_prediction_oneoff_regression = _plugin_loader.list(plugin_type="metric")["prediction"]["one_off"][
+    "regression"
 ]
-"""Evaluation metrics supported in the regression task setting.
-
-Possible values:
-    - ``"r2"``:
-        R^2 (coefficient of determination) regression score function.
-    - ``"mse"``:
-        Mean squared error regression loss.
-    - ``"mae"``:
-        Mean absolute error regression loss.
-"""
+builtin_metrics_time_to_event = _plugin_loader.list(plugin_type="metric")["time_to_event"]
 
-TimeToEventSupportedMetric = Literal[
-    "c_index",
-    "brier_score",
-]
-"""Evaluation metrics supported in the time-to-event (survival) task setting.
-
-Possible values:
-    - ``"c_index"``:
-        Concordance index based on inverse probability of censoring weights.
-    - ``"brier_score"``:
-        The time-dependent Brier score.
-"""
 
 OutputMetric = Literal[
     "min",
@@ -157,15 +72,6 @@
         Average duration for the fold evaluation.
 """
 
-classifier_supported_metrics = get_args(ClassifierSupportedMetric)
-"""A tuple of all possible values of :obj:`~tempor.benchmarks.evaluation.ClassifierSupportedMetric`."""
-
-regression_supported_metrics = get_args(RegressionSupportedMetric)
-"""A tuple of all possible values of :obj:`~tempor.benchmarks.evaluation.RegressionSupportedMetric`."""
-
-time_to_event_supported_metrics = get_args(TimeToEventSupportedMetric)
-"""A tuple of all possible values of :obj:`~tempor.benchmarks.evaluation.TimeToEventSupportedMetric`."""
-
 output_metrics = get_args(OutputMetric)
 """A tuple of all possible values of :obj:`~tempor.benchmarks.evaluation.OutputMetric`."""
 
@@ -220,152 +126,6 @@ def _postprocess_results(results: _InternalScores) -> pd.DataFrame:
     return output
 
 
-class ClassifierMetrics:
-    @pydantic_utils.validate_arguments
-    def __init__(
-        self,
-        metric: Union[ClassifierSupportedMetric, Sequence[ClassifierSupportedMetric]] = classifier_supported_metrics,
-    ) -> None:
-        """Helper class for evaluating the performance of the classifier.
-
-        Args:
-            metric (Union[ClassifierSupportedMetric, Sequence[ClassifierSupportedMetric]], optional):
-                The type of metric(s) to use for evaluation.
-                A string (one of :obj:`~tempor.benchmarks.evaluation.ClassifierSupportedMetric`) or a sequence of such.
-                Defaults to :obj:`~tempor.benchmarks.evaluation.classifier_supported_metrics`.
-        """
-        self.metrics: Union[ClassifierSupportedMetric, Sequence[ClassifierSupportedMetric]]
-        if isinstance(metric, str):  # pragma: no cover
-            # Should be prevented by pydantic.
-            self.metrics = [cast(ClassifierSupportedMetric, metric)]
-        else:
-            self.metrics = metric
-
-    def score_proba(self, y_test: np.ndarray, y_pred_proba: np.ndarray) -> Dict[str, float]:
-        """Return the evaluation metrics for the given predictions, where the predictions are of
-        predicted probabilities form.
-
-        Args:
-            y_test (np.ndarray): Test labels.
-            y_pred_proba (np.ndarray): Predicted probabilities.
-
-        Returns:
-            Dict[str, float]: Evaluation score.
-        """
-        if y_test is None or y_pred_proba is None:
-            raise ValueError("Invalid input for score_proba")
-
-        results = {}
-        y_pred = np.argmax(np.asarray(y_pred_proba), axis=1)
-        for metric in self.metrics:
-            if metric == "aucprc":
-                results[metric] = self.average_precision_score(y_test, y_pred_proba)
-            elif metric == "aucroc":
-                results[metric] = self.roc_auc_score(y_test, y_pred_proba)
-            elif metric == "accuracy":
-                results[metric] = sklearn.metrics.accuracy_score(y_test, y_pred)
-            elif metric == "f1_score_micro":
-                results[metric] = sklearn.metrics.f1_score(
-                    y_test,
-                    y_pred,
-                    average="micro",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "f1_score_macro":
-                results[metric] = sklearn.metrics.f1_score(
-                    y_test,
-                    y_pred,
-                    average="macro",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "f1_score_weighted":
-                results[metric] = sklearn.metrics.f1_score(
-                    y_test,
-                    y_pred,
-                    average="weighted",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "kappa":
-                results[metric] = sklearn.metrics.cohen_kappa_score(y_test, y_pred)
-            elif metric == "kappa_quadratic":
-                results[metric] = sklearn.metrics.cohen_kappa_score(y_test, y_pred, weights="quadratic")
-            elif metric == "recall_micro":
-                results[metric] = sklearn.metrics.recall_score(
-                    y_test,
-                    y_pred,
-                    average="micro",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "recall_macro":
-                results[metric] = sklearn.metrics.recall_score(
-                    y_test,
-                    y_pred,
-                    average="macro",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "recall_weighted":
-                results[metric] = sklearn.metrics.recall_score(
-                    y_test,
-                    y_pred,
-                    average="weighted",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "precision_micro":
-                results[metric] = sklearn.metrics.precision_score(
-                    y_test,
-                    y_pred,
-                    average="micro",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "precision_macro":
-                results[metric] = sklearn.metrics.precision_score(
-                    y_test,
-                    y_pred,
-                    average="macro",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "precision_weighted":
-                results[metric] = sklearn.metrics.precision_score(
-                    y_test,
-                    y_pred,
-                    average="weighted",
-                    zero_division=0,  # pyright: ignore
-                )
-            elif metric == "mcc":
-                results[metric] = sklearn.metrics.matthews_corrcoef(y_test, y_pred)
-            else:  # pragma: no cover
-                raise ValueError(f"invalid metric {metric}")
-
-        logger.debug(f"evaluate_classifier: {results}")
-        return results
-
-    def roc_auc_score(self, y_test: np.ndarray, y_pred_proba: np.ndarray) -> float:
-        """Return the ROC AUC score for the given predictions, where the predictions are of
-        predicted probabilities form.
-
-        Args:
-            y_test (np.ndarray): Test labels.
-            y_pred_proba (np.ndarray): Predicted probabilities.
-
-        Returns:
-            float: ROC AUC score.
-        """
-        return utils.evaluate_auc_multiclass(y_test, y_pred_proba)[0]
-
-    def average_precision_score(self, y_test: np.ndarray, y_pred_proba: np.ndarray) -> float:
-        """Return the average precision score for the given predictions, where the predictions are of
-        predicted probabilities form.
-
-        Args:
-            y_test (np.ndarray): Test labels.
-            y_pred_proba (np.ndarray): Predicted probabilities.
-
-        Returns:
-            float: Average precision score.
-        """
-        return utils.evaluate_auc_multiclass(y_test, y_pred_proba)[1]
-
-
 @pydantic_utils.validate_arguments(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
 def evaluate_prediction_oneoff_classifier(  # pylint: disable=unused-argument
     estimator: Any,
@@ -403,10 +163,18 @@ def evaluate_prediction_oneoff_classifier(  # pylint: disable=unused-argument
             The columns of the dataframe contain details about the cross-validation repeats: one column for each
             :obj:`~tempor.benchmarks.evaluation.OutputMetric`.
 
-            The index of the dataframe contains all the metrics evaluated: all of
-            :obj:`~tempor.benchmarks.evaluation.ClassifierSupportedMetric`.
+            The index of the dataframe contains all the metrics registered:
+            >>> from tempor import plugin_loader
+            >>> plugin_loader.list(plugin_type="metric")["prediction"]["one_off"]["classification"]
+            [...]
+
     """
 
+    # For the sake of import modularity, do not use the global plugin loader here, but create own:
+    _plugin_loader = plugins.PluginLoader()
+    metric_plugin_category = "prediction.one_off.classification"
+    supported_metrics = _plugin_loader.list(plugin_type="metric")["prediction"]["one_off"]["classification"]
+
     with warnings.catch_warnings():
         if silence_warnings:
             warnings.simplefilter("ignore")
@@ -417,9 +185,8 @@ def evaluate_prediction_oneoff_classifier(  # pylint: disable=unused-argument
         enable_reproducibility(random_state)
 
         results = _InternalScores()
-        evaluator = ClassifierMetrics()
-        for metric in classifier_supported_metrics:
-            results.metrics[metric] = np.zeros(n_splits)
+        for metric_name in supported_metrics:
+            results.metrics[metric_name] = np.zeros(n_splits)
 
         splitter = sklearn.model_selection.StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
 
@@ -433,6 +200,7 @@ def evaluate_prediction_oneoff_classifier(  # pylint: disable=unused-argument
         for train_data, test_data in data.split(splitter=splitter, y=labels):
             model = copy.deepcopy(estimator_)
             start = time()
+
             try:
                 model.fit(train_data)
 
@@ -441,10 +209,15 @@ def evaluate_prediction_oneoff_classifier(  # pylint: disable=unused-argument
                 test_labels = test_data.predictive.targets.numpy()
                 preds = model.predict_proba(test_data).numpy()
 
-                scores = evaluator.score_proba(test_labels, preds)
-                for metric in scores:
-                    results.metrics[metric][indx] = scores[metric]
+                for metric_name in supported_metrics:
+                    metric = cast(
+                        metric_module.OneOffClassificationMetric,
+                        _plugin_loader.get(f"{metric_plugin_category}.{metric_name}", plugin_type="metric"),
+                    )
+                    results.metrics[metric_name][indx] = metric.evaluate(test_labels, preds)
+
                 results.errors.append(0)
+
             except BaseException as e:  # pylint: disable=broad-except
                 logger.error(f"Evaluation failed: {e}")
                 results.errors.append(1)
@@ -494,10 +267,18 @@ def evaluate_prediction_oneoff_regressor(  # pylint: disable=unused-argument
             The columns of the dataframe contain details about the cross-validation repeats: one column for each
             :obj:`~tempor.benchmarks.evaluation.OutputMetric`.
 
-            The index of the dataframe contains all the metrics evaluated: all of
-            :obj:`~tempor.benchmarks.evaluation.RegressionSupportedMetric`.
+            The index of the dataframe contains all the metrics registered:
+            >>> from tempor import plugin_loader
+            >>> plugin_loader.list(plugin_type="metric")["prediction"]["one_off"]["regression"]
+            [...]
+
     """
 
+    # For the sake of import modularity, do not use the global plugin loader here, but create own:
+    _plugin_loader = plugins.PluginLoader()
+    metric_plugin_category = "prediction.one_off.regression"
+    supported_metrics = _plugin_loader.list(plugin_type="metric")["prediction"]["one_off"]["regression"]
+
     with warnings.catch_warnings():
         if silence_warnings:
             warnings.simplefilter("ignore")
@@ -506,10 +287,9 @@ def evaluate_prediction_oneoff_regressor(  # pylint: disable=unused-argument
             raise ValueError("n_splits must be an integer >= 2")
         estimator_ = cast("BaseOneOffRegressor", estimator)
         enable_reproducibility(random_state)
-        metrics = regression_supported_metrics
 
         results = _InternalScores()
-        for metric in metrics:
+        for metric in supported_metrics:
             results.metrics[metric] = np.zeros(n_splits)
 
         splitter = sklearn.model_selection.KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
@@ -518,6 +298,7 @@ def evaluate_prediction_oneoff_regressor(  # pylint: disable=unused-argument
         for train_data, test_data in data.split(splitter=splitter):
             model = copy.deepcopy(estimator_)
             start = time()
+
             try:
                 model.fit(train_data)
 
@@ -526,10 +307,15 @@ def evaluate_prediction_oneoff_regressor(  # pylint: disable=unused-argument
                 targets = test_data.predictive.targets.numpy().squeeze()
                 preds = model.predict(test_data).numpy().squeeze()
 
-                results.metrics["mse"][indx] = sklearn.metrics.mean_squared_error(targets, preds)
-                results.metrics["mae"][indx] = sklearn.metrics.mean_absolute_error(targets, preds)
-                results.metrics["r2"][indx] = sklearn.metrics.r2_score(targets, preds)
+                for metric_name in supported_metrics:
+                    metric = cast(
+                        metric_module.OneOffRegressionMetric,
+                        _plugin_loader.get(f"{metric_plugin_category}.{metric_name}", plugin_type="metric"),
+                    )
+                    results.metrics[metric_name][indx] = metric.evaluate(targets, preds)
+
                 results.errors.append(0)
+
             except BaseException as e:  # pylint: disable=broad-except
                 logger.error(f"Regression evaluation failed: {e}")
                 results.errors.append(1)
@@ -542,72 +328,17 @@ def evaluate_prediction_oneoff_regressor(  # pylint: disable=unused-argument
     return _postprocess_results(results)
 
 
-TimeToEventMetricCallable = Callable[[np.ndarray, np.ndarray, np.ndarray, List[float]], List[float]]
-"""Standardized function for time-to-event metric.
-
-Inputs are:
-    * ``training_array_struct`` (np.ndarray)
-    * ``testing_array_struct`` (np.ndarray)
-    * ``predictions`` (np.ndarray)
-    * ``horizons`` (List[float])
-
-Output is:
-    A list with the metric values for each horizon.
-"""
-
-
-def compute_c_index(
-    training_array_struct: np.ndarray, testing_array_struct: np.ndarray, predictions: np.ndarray, horizons: List[float]
-) -> List[float]:
-    """Compute the IPCW concordance index.
-
-    Args:
-        training_array_struct (np.ndarray): Training data as a structured array.
-        testing_array_struct (np.ndarray): Testing data as a structured array.
-        predictions (np.ndarray): Predictions.
-        horizons (List[float]): Evaluation horizons.
-
-    Returns:
-        List[float]: List of metric values for each horizon.
-    """
-    metrics: List[float] = []
-    for horizon_idx, horizon_time in enumerate(horizons):
-        predictions_at_horizon_time = predictions[:, horizon_idx, :].reshape((-1,))
-        out = tempor_metrics.concordance_index_ipcw(
-            training_array_struct, testing_array_struct, predictions_at_horizon_time, float(horizon_time)
-        )
-        metrics.append(out[0])
-    return metrics
-
-
-def compute_brier_score(
-    training_array_struct: np.ndarray, testing_array_struct: np.ndarray, predictions: np.ndarray, horizons: List[float]
-) -> List[float]:
-    """Compute the time-dependent Brier score.
-
-    Args:
-        training_array_struct (np.ndarray): Training data as a structured array.
-        testing_array_struct (np.ndarray): Testing data as a structured array.
-        predictions (np.ndarray): Predictions.
-        horizons (List[float]): Evaluation horizons.
-
-    Returns:
-        List[float]: List of metric values for each horizon.
-    """
-    predictions = predictions.reshape((predictions.shape[0], predictions.shape[1]))
-    times, scores = tempor_metrics.brier_score(  # pylint: disable=unused-variable
-        training_array_struct, testing_array_struct, predictions, horizons
-    )
-    return scores.tolist()
-
-
-def _compute_time_to_event_metric(
-    metric_func: TimeToEventMetricCallable,
+def _prep_data_for_time_to_event_metric(
     train_data: dataset.TimeToEventAnalysisDataset,
     test_data: dataset.TimeToEventAnalysisDataset,
     horizons: data_typing.TimeIndex,
     predictions: samples.TimeSeriesSamples,
-) -> float:
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[float], np.ndarray, np.ndarray]:
+    """Validate the data for time-to-event metric evaluation and
+    prepare the data in the format expected by the metric evaluation call.
+    """
+
+    # Validate the data.
     if not predictions.num_timesteps_equal():
         raise ValueError(
             f"Expected time to event prediction values for horizons {horizons} all to have equal number of time steps "
@@ -628,18 +359,15 @@ def _compute_time_to_event_metric(
         float(horizons[0])  # pyright: ignore
     except (TypeError, ValueError) as e:
         raise ValueError("Currently only int or float time horizons supported.") from e
+
+    # Prepare the data.
     horizons = cast(List[float], horizons)
 
     predictions_array = predictions.numpy()
     t_train, y_train = (df.to_numpy().reshape((-1,)) for df in train_data.predictive.targets.split_as_two_dataframes())
-    y_train_struct = tempor_metrics.create_structured_array(y_train, t_train)
     t_test, y_test = (df.to_numpy().reshape((-1,)) for df in test_data.predictive.targets.split_as_two_dataframes())
-    y_test_struct = tempor_metrics.create_structured_array(y_test, t_test)
 
-    metrics: List[float] = metric_func(y_train_struct, y_test_struct, predictions_array, horizons)
-    avg_metric = float(np.asarray(metrics).mean())
-
-    return avg_metric
+    return y_test, t_test, predictions_array, horizons, y_train, t_train
 
 
 @pydantic_utils.validate_arguments(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
@@ -682,9 +410,18 @@ def evaluate_time_to_event(  # pylint: disable=unused-argument
             The columns of the dataframe contain details about the cross-validation repeats: one column for each
             :obj:`~tempor.benchmarks.evaluation.OutputMetric`.
 
-            The index of the dataframe contains all the metrics evaluated: all of
-            :obj:`~tempor.benchmarks.evaluation.TimeToEventSupportedMetric`.
+            The index of the dataframe contains all the metrics registered:
+            >>> from tempor import plugin_loader
+            >>> plugin_loader.list(plugin_type="metric")["time_to_event"]
+            [...]
+
     """
+
+    # For the sake of import modularity, do not use the global plugin loader here, but create own:
+    _plugin_loader = plugins.PluginLoader()
+    metric_plugin_category = "time_to_event"
+    supported_metrics = _plugin_loader.list(plugin_type="metric")["time_to_event"]
+
     with warnings.catch_warnings():
         if silence_warnings:
             warnings.simplefilter("ignore")
@@ -694,14 +431,9 @@ def evaluate_time_to_event(  # pylint: disable=unused-argument
             raise ValueError("n_splits must be an integer >= 2")
         estimator_ = cast("BaseTimeToEventAnalysis", estimator)
         enable_reproducibility(random_state)
-        metrics = time_to_event_supported_metrics
-        metrics_map = {
-            "c_index": compute_c_index,
-            "brier_score": compute_brier_score,
-        }
 
         results = _InternalScores()
-        for metric in metrics:
+        for metric in supported_metrics:
             results.metrics[metric] = np.zeros(n_splits)
 
         splitter = sklearn.model_selection.KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
@@ -716,15 +448,22 @@ def evaluate_time_to_event(  # pylint: disable=unused-argument
                 # targets = test_data.predictive.targets.numpy().squeeze()
                 preds = model.predict(test_data, horizons=horizons)
 
-                for metric_name in time_to_event_supported_metrics:
-                    metric_func = metrics_map[metric_name]
-                    results.metrics[metric_name][indx] = _compute_time_to_event_metric(
-                        metric_func,
-                        train_data=train_data,
-                        test_data=test_data,
-                        horizons=horizons,
-                        predictions=preds,
+                y_test, t_test, predictions_array, horizons, y_train, t_train = _prep_data_for_time_to_event_metric(
+                    train_data=train_data,
+                    test_data=test_data,
+                    horizons=horizons,
+                    predictions=preds,
+                )
+                for metric_name in supported_metrics:
+                    metric = cast(
+                        metric_module.TimeToEventMetric,
+                        _plugin_loader.get(f"{metric_plugin_category}.{metric_name}", plugin_type="metric"),
+                    )
+                    metric_per_horizon = metric.evaluate(
+                        (y_test, t_test), predictions_array, horizons, (y_train, t_train)
                     )
+                    avg_metric = float(np.asarray(metric_per_horizon).mean())
+                    results.metrics[metric_name][indx] = avg_metric
 
                 results.errors.append(0)
 
diff --git a/src/tempor/benchmarks/utils.py b/src/tempor/benchmarks/utils.py
index e4bf5931..e7489ac6 100644
--- a/src/tempor/benchmarks/utils.py
+++ b/src/tempor/benchmarks/utils.py
@@ -1,80 +1,8 @@
 """Any utilities for the ``benchmark`` package directory."""
 
-from typing import Tuple, cast
+from typing import Tuple
 
 import numpy as np
-from sklearn.metrics import auc, average_precision_score, precision_recall_curve, roc_auc_score, roc_curve
-from sklearn.preprocessing import label_binarize
-
-from tempor.log import logger as log
-
-
-def get_y_pred_proba_hlpr(y_pred_proba: np.ndarray, nclasses: int) -> np.ndarray:
-    """A helper utility for inferring the correct y_pred_proba for multiclass situations, specifically in the case
-    of binary classification. See source code for the specifics.
-
-    Args:
-        y_pred_proba (np.ndarray): Predicted probabilities.
-        nclasses (int): Number of classes.
-
-    Returns:
-        np.ndarray: The correctly inferred ``y_pred_proba``.
-    """
-    if nclasses == 2:
-        if len(y_pred_proba.shape) < 2:
-            return y_pred_proba
-
-        if y_pred_proba.shape[1] == 2:
-            return y_pred_proba[:, 1]
-
-    return y_pred_proba
-
-
-def evaluate_auc_multiclass(
-    y_test: np.ndarray,
-    y_pred_proba: np.ndarray,
-) -> Tuple[float, float]:
-    """Helper for evaluating AUCROC/AUCPRC for any number of classes."""
-
-    y_test = np.asarray(y_test)
-    y_pred_proba = np.asarray(y_pred_proba)
-
-    nnan = sum(np.ravel(np.isnan(y_pred_proba)))
-
-    if nnan:
-        raise ValueError("nan in predictions. aborting")
-
-    n_classes = len(set(np.ravel(y_test)))
-    classes = sorted(set(np.ravel(y_test)))
-
-    y_pred_proba_tmp = get_y_pred_proba_hlpr(y_pred_proba, n_classes)
-
-    if n_classes > 2:
-        log.debug(f"+evaluate_auc {y_test.shape} {y_pred_proba_tmp.shape}")
-
-        fpr = dict()
-        tpr = dict()
-        precision = dict()
-        recall = dict()
-        average_precision = dict()
-        roc_auc: dict = dict()
-
-        y_test = cast(np.ndarray, label_binarize(y_test, classes=classes, sparse_output=False))
-
-        fpr["micro"], tpr["micro"], _ = roc_curve(y_test.ravel(), y_pred_proba_tmp.ravel())
-        roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])
-
-        precision["micro"], recall["micro"], _ = precision_recall_curve(y_test.ravel(), y_pred_proba_tmp.ravel())
-
-        average_precision["micro"] = average_precision_score(y_test, y_pred_proba_tmp, average="micro")
-
-        aucroc = roc_auc["micro"]
-        aucprc = average_precision["micro"]
-    else:
-        aucroc = roc_auc_score(np.ravel(y_test), y_pred_proba_tmp, multi_class="ovr")
-        aucprc = average_precision_score(np.ravel(y_test), y_pred_proba_tmp)
-
-    return float(aucroc), float(aucprc)
 
 
 def generate_score(metric: np.ndarray) -> Tuple[float, float]:
diff --git a/src/tempor/core/utils.py b/src/tempor/core/utils.py
index 93f1a84f..4d2e6ed8 100644
--- a/src/tempor/core/utils.py
+++ b/src/tempor/core/utils.py
@@ -179,3 +179,54 @@ def unique_in_order_of_appearance(iterable: Iterable) -> List:
         List: List of unique elements in order of their appearance.
     """
     return list(dict.fromkeys(iterable))
+
+
+def is_method_defined_in_class(cls_or_obj: Any, method_name: str) -> bool:
+    """Check if method named ``method_name`` method is defined in the given class (or object's class) or inherited.
+
+    Args:
+        cls_or_obj (Any): The class or object to check.
+        method_name (str): The name of the method to check.
+
+    Returns:
+        bool: True if method is defined in ``cls``, `False` if inherited.
+    """
+    init_qualname = getattr(cls_or_obj, method_name).__qualname__
+    class_name = cls_or_obj.__name__ if isinstance(cls_or_obj, type) else cls_or_obj.__class__.__name__
+    return init_qualname.startswith(class_name + ".")
+
+
+def clean_multiline_docstr(docstr: str) -> str:
+    """Clean a multi-line docstring by getting rid of newlines and cleaning up whitespace.
+
+    Args:
+        docstr (str): The docstring to clean.
+
+    Returns:
+        str: The cleaned docstring.
+    """
+    return " ".join([line.strip() for line in docstr.split("\n")]).strip()
+
+
+def make_description_from_doc(obj: Any, max_len_keep: int = 100) -> str:
+    """Make a description from the docstring of an object. Take the class docstring and the ``__init__`` docstring
+    (if there was one defined on the class). If the combined length of these is greater than ``max_len_keep``, then
+    truncate with ``...``.
+
+    Args:
+        obj (Any): The object to get the description of.
+        max_len_keep (int, optional): Maximum description length before truncating. Defaults to ``100``.
+
+    Returns:
+        str: Description of the object.
+    """
+    class_doc: str = obj.__doc__ if obj.__doc__ is not None else ""
+    init_doc: str = ""
+    if is_method_defined_in_class(obj, "__init__"):
+        init_doc = obj.__init__.__doc__ if obj.__init__.__doc__ is not None else ""
+    class_and_init_docs_combo = f"{class_doc} {init_doc}".strip()
+
+    if len(class_and_init_docs_combo) > max_len_keep:
+        return class_and_init_docs_combo[:max_len_keep] + "..."
+    else:
+        return class_and_init_docs_combo
diff --git a/src/tempor/data/utils.py b/src/tempor/data/utils.py
index 11f18840..f5e4c1c3 100644
--- a/src/tempor/data/utils.py
+++ b/src/tempor/data/utils.py
@@ -322,7 +322,7 @@ def make_sample_time_index_tuples(
     """
     if len(sample_index) != len(time_indexes):
         raise ValueError("Expected the same number of elements in `sample_index` and `time_indexes`")
-    sample_indexes_copied = [[si] * len(tis) for si, tis in zip(sample_index, time_indexes)]  # type: ignore[arg-type]
+    sample_indexes_copied = [[si] * len(tis) for si, tis in zip(sample_index, time_indexes)]
     sample_indexes_flattened = list(itertools.chain.from_iterable(sample_indexes_copied))
     time_indexes_flattened = list(itertools.chain.from_iterable(time_indexes))
     pairs = [(si, ti) for si, ti in zip(sample_indexes_flattened, time_indexes_flattened)]
diff --git a/src/tempor/datasources/datasource.py b/src/tempor/datasources/datasource.py
index 048e134f..4a05b8f1 100644
--- a/src/tempor/datasources/datasource.py
+++ b/src/tempor/datasources/datasource.py
@@ -3,8 +3,10 @@
 # pylint: disable=unnecessary-ellipsis
 
 import abc
+import contextlib
 import os
-from typing import Any, ClassVar, Optional, Type
+import ssl
+from typing import Any, ClassVar, Generator, Optional, Type
 
 import tempor
 from tempor.core import plugins
@@ -17,6 +19,23 @@
 """
 
 
+@contextlib.contextmanager
+def monkeypatch_ssl_error_workaround() -> Generator:  # pragma: no cover
+    """Some datasets (e.g. UCI diabetes) are hosted on servers that may have SSL issues. This is a workaround that
+    monkeypatches the `ssl` module to ignore SSL errors temporarily.
+    """
+
+    context_backup = ssl._create_default_https_context  # pylint: disable=protected-access
+    ssl._create_default_https_context = ssl._create_unverified_context  # pylint: disable=protected-access
+
+    try:
+        yield
+
+    finally:
+        # Make sure to restore the original secure context!
+        ssl._create_default_https_context = context_backup  # pylint: disable=protected-access
+
+
 # TODO: Unit test.
 
 
@@ -123,6 +142,9 @@ def load(self, **kwargs: Any) -> dataset.OneOffPredictionDataset:  # pragma: no
         ...
 
 
+plugins.register_plugin_category("prediction.one_off", OneOffPredictionDataSource, plugin_type="datasource")
+
+
 class TemporalPredictionDataSource(DataSource):
     @property
     def predictive_task(self) -> data_typing.PredictiveTask:
@@ -146,6 +168,9 @@ def load(self, **kwargs: Any) -> dataset.TemporalPredictionDataset:  # pragma: n
         ...
 
 
+plugins.register_plugin_category("prediction.temporal", TemporalPredictionDataSource, plugin_type="datasource")
+
+
 class TimeToEventAnalysisDataSource(DataSource):
     @property
     def predictive_task(self) -> data_typing.PredictiveTask:
@@ -169,6 +194,9 @@ def load(self, **kwargs: Any) -> dataset.TimeToEventAnalysisDataset:  # pragma:
         ...
 
 
+plugins.register_plugin_category("time_to_event", TimeToEventAnalysisDataSource, plugin_type="datasource")
+
+
 class OneOffTreatmentEffectsDataSource(DataSource):
     @property
     def predictive_task(self) -> data_typing.PredictiveTask:
@@ -192,6 +220,9 @@ def load(self, **kwargs: Any) -> dataset.OneOffTreatmentEffectsDataset:  # pragm
         ...
 
 
+plugins.register_plugin_category("treatments.one_off", OneOffTreatmentEffectsDataSource, plugin_type="datasource")
+
+
 class TemporalTreatmentEffectsDataSource(DataSource):
     @property
     def predictive_task(self) -> data_typing.PredictiveTask:
@@ -213,3 +244,6 @@ def load(self, **kwargs: Any) -> dataset.TemporalTreatmentEffectsDataset:  # pra
             dataset.TemporalTreatmentEffectsDataset: The loaded dataset.
         """
         ...
+
+
+plugins.register_plugin_category("treatments.temporal", TemporalTreatmentEffectsDataSource, plugin_type="datasource")
diff --git a/src/tempor/datasources/prediction/one_off/__init__.py b/src/tempor/datasources/prediction/one_off/__init__.py
index 54649c3d..c87249b9 100644
--- a/src/tempor/datasources/prediction/one_off/__init__.py
+++ b/src/tempor/datasources/prediction/one_off/__init__.py
@@ -1,9 +1,6 @@
 """Package directory for one-off prediction data sources."""
 
 from tempor.core import plugins
-from tempor.datasources import datasource
-
-plugins.register_plugin_category("prediction.one_off", datasource.OneOffPredictionDataSource, plugin_type="datasource")
 
 plugins.importing.import_plugins(__file__)
 
diff --git a/src/tempor/datasources/prediction/temporal/__init__.py b/src/tempor/datasources/prediction/temporal/__init__.py
index 31f40df6..5ce92efc 100644
--- a/src/tempor/datasources/prediction/temporal/__init__.py
+++ b/src/tempor/datasources/prediction/temporal/__init__.py
@@ -1,11 +1,6 @@
 """Package for temporal prediction data sources."""
 
 from tempor.core import plugins
-from tempor.datasources import datasource
-
-plugins.register_plugin_category(
-    "prediction.temporal", datasource.TemporalPredictionDataSource, plugin_type="datasource"
-)
 
 plugins.importing.import_plugins(__file__)
 
diff --git a/src/tempor/datasources/prediction/temporal/plugin_uci_diabetes.py b/src/tempor/datasources/prediction/temporal/plugin_uci_diabetes.py
index 26397248..5aec4dca 100644
--- a/src/tempor/datasources/prediction/temporal/plugin_uci_diabetes.py
+++ b/src/tempor/datasources/prediction/temporal/plugin_uci_diabetes.py
@@ -16,6 +16,8 @@
 from tempor.datasources import datasource
 from tempor.log import logger
 
+from ...datasource import monkeypatch_ssl_error_workaround
+
 
 @plugins.register_plugin(name="uci_diabetes", category="prediction.temporal", plugin_type="datasource")
 class UCIDiabetesDataSource(datasource.TemporalPredictionDataSource):
@@ -53,6 +55,15 @@ def url() -> str:  # noqa: D102
     def dataset_dir() -> str:  # noqa: D102
         return str(Path(UCIDiabetesDataSource.data_root_dir) / "uci_diabetes")
 
+    def _load(self) -> Any:
+        return uci_diabetes(
+            data_home=UCIDiabetesDataSource.data_root_dir,
+            refresh_cache=True,
+            redownload=False,
+            make_regular=self.make_regular,
+            use_int_index=self.use_int_index,
+        )
+
     def load(self, **kwargs: Any) -> dataset.TemporalPredictionDataset:  # noqa: D102
         download_retries = 3
         download_pause_sec = 5
@@ -61,14 +72,17 @@ def load(self, **kwargs: Any) -> dataset.TemporalPredictionDataset:  # noqa: D10
             # TODO: May wish to exclude this from tests / download files from a more stable location /
             # make tests resilient to internet connection failures.
             try:
-                clv_dataset = uci_diabetes(
-                    data_home=UCIDiabetesDataSource.data_root_dir,
-                    refresh_cache=True,
-                    redownload=False,
-                    make_regular=self.make_regular,
-                    use_int_index=self.use_int_index,
-                )
-            except (requests.exceptions.RequestException, urllib.error.URLError) as ex:  # pragma: no cover
+                clv_dataset = self._load()
+            except urllib.error.URLError as ex:  # pragma: no cover
+                if "SSL" in str(ex) or "ssl" in str(ex):
+                    with monkeypatch_ssl_error_workaround():
+                        clv_dataset = self._load()
+                else:
+                    raise
+            except (  # pylint: disable=duplicate-except
+                requests.exceptions.RequestException,
+                urllib.error.URLError,
+            ) as ex:  # pragma: no cover
                 if retry + 1 == download_retries:
                     logger.error(f"Failed to download UCI diabetes dataset after {download_retries} retries.")
                     raise
diff --git a/src/tempor/datasources/time_to_event/__init__.py b/src/tempor/datasources/time_to_event/__init__.py
index aa3679e7..f819a58a 100644
--- a/src/tempor/datasources/time_to_event/__init__.py
+++ b/src/tempor/datasources/time_to_event/__init__.py
@@ -1,9 +1,6 @@
 """Package for time to event analysis data sources."""
 
 from tempor.core import plugins
-from tempor.datasources import datasource
-
-plugins.register_plugin_category("time_to_event", datasource.TimeToEventAnalysisDataSource, plugin_type="datasource")
 
 plugins.importing.import_plugins(__file__)
 
diff --git a/src/tempor/datasources/treatments/one_off/__init__.py b/src/tempor/datasources/treatments/one_off/__init__.py
index 7e9b6caf..83ab1cd9 100644
--- a/src/tempor/datasources/treatments/one_off/__init__.py
+++ b/src/tempor/datasources/treatments/one_off/__init__.py
@@ -1,11 +1,6 @@
 """Package for one-off treatment effects data sources."""
 
 from tempor.core import plugins
-from tempor.datasources import datasource
-
-plugins.register_plugin_category(
-    "treatments.one_off", datasource.OneOffTreatmentEffectsDataSource, plugin_type="datasource"
-)
 
 plugins.importing.import_plugins(__file__)
 
diff --git a/src/tempor/datasources/treatments/temporal/__init__.py b/src/tempor/datasources/treatments/temporal/__init__.py
index 82ffd34d..460e388d 100644
--- a/src/tempor/datasources/treatments/temporal/__init__.py
+++ b/src/tempor/datasources/treatments/temporal/__init__.py
@@ -1,11 +1,6 @@
 """Package for temporal treatment effects data sources."""
 
 from tempor.core import plugins
-from tempor.datasources import datasource
-
-plugins.register_plugin_category(
-    "treatments.temporal", datasource.TemporalTreatmentEffectsDataSource, plugin_type="datasource"
-)
 
 plugins.importing.import_plugins(__file__)
 
diff --git a/src/tempor/methods/prediction/temporal/classification/__init__.py b/src/tempor/methods/prediction/temporal/classification/__init__.py
index 70faca48..e367f196 100644
--- a/src/tempor/methods/prediction/temporal/classification/__init__.py
+++ b/src/tempor/methods/prediction/temporal/classification/__init__.py
@@ -43,7 +43,7 @@ def fit(self, data: dataset.BaseDataset, *args: Any, **kwargs: Any) -> Self:  #
         return self
 
     @pydantic_utils.validate_arguments(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
-    def predict(  # type: ignore [override]  # pylint: disable=arguments-differ
+    def predict(  # pylint: disable=arguments-differ
         self,
         data: dataset.PredictiveDataset,
         n_future_steps: int,
@@ -55,7 +55,7 @@ def predict(  # type: ignore [override]  # pylint: disable=arguments-differ
         return super().predict(data, n_future_steps, *args, time_delta=time_delta, **kwargs)
 
     @pydantic_utils.validate_arguments(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
-    def predict_proba(  # type: ignore [override]  # pylint: disable=arguments-differ
+    def predict_proba(  # pylint: disable=arguments-differ
         self,
         data: dataset.PredictiveDataset,
         n_future_steps: int,
@@ -67,13 +67,13 @@ def predict_proba(  # type: ignore [override]  # pylint: disable=arguments-diffe
         return super().predict_proba(data, n_future_steps, *args, time_delta=time_delta, **kwargs)
 
     @abc.abstractmethod
-    def _predict(  # type: ignore[override]  # pylint: disable=arguments-differ
+    def _predict(  # pylint: disable=arguments-differ
         self, data: dataset.PredictiveDataset, n_future_steps: int, *args: Any, time_delta: int = 1, **kwargs: Any
     ) -> samples.TimeSeriesSamples:  # pragma: no cover
         ...
 
     @abc.abstractmethod
-    def _predict_proba(  # type: ignore[override]  # pylint: disable=arguments-differ
+    def _predict_proba(  # pylint: disable=arguments-differ
         self, data: dataset.PredictiveDataset, n_future_steps: int, *args: Any, time_delta: int = 1, **kwargs: Any
     ) -> samples.TimeSeriesSamples:  # pragma: no cover
         ...
diff --git a/src/tempor/methods/prediction/temporal/classification/plugin_seq2seq_classifier.py b/src/tempor/methods/prediction/temporal/classification/plugin_seq2seq_classifier.py
index 8bef72e7..74d0e127 100644
--- a/src/tempor/methods/prediction/temporal/classification/plugin_seq2seq_classifier.py
+++ b/src/tempor/methods/prediction/temporal/classification/plugin_seq2seq_classifier.py
@@ -69,7 +69,7 @@ def _fit(
         self.model.fit(cl_dataset)
         return self
 
-    def _predict(  # type: ignore[override]
+    def _predict(
         self,
         data: dataset.PredictiveDataset,
         n_future_steps: int,
diff --git a/src/tempor/methods/prediction/temporal/regression/__init__.py b/src/tempor/methods/prediction/temporal/regression/__init__.py
index f8f6271e..3e269c7e 100644
--- a/src/tempor/methods/prediction/temporal/regression/__init__.py
+++ b/src/tempor/methods/prediction/temporal/regression/__init__.py
@@ -43,7 +43,7 @@ def fit(self, data: dataset.BaseDataset, *args: Any, **kwargs: Any) -> Self:  #
         return self
 
     @pydantic_utils.validate_arguments(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
-    def predict(  # type: ignore [override] # pylint: disable=arguments-differ
+    def predict(  # pylint: disable=arguments-differ
         self,
         data: dataset.PredictiveDataset,
         n_future_steps: int,
@@ -55,7 +55,7 @@ def predict(  # type: ignore [override] # pylint: disable=arguments-differ
         return super().predict(data, n_future_steps, *args, time_delta=time_delta, **kwargs)
 
     @abc.abstractmethod
-    def _predict(  # type: ignore[override]  # pylint: disable=arguments-differ
+    def _predict(  # pylint: disable=arguments-differ
         self, data: dataset.PredictiveDataset, n_future_steps: int, *args: Any, time_delta: int = 1, **kwargs: Any
     ) -> samples.TimeSeriesSamples:  # pragma: no cover
         ...
diff --git a/src/tempor/methods/prediction/temporal/regression/plugin_seq2seq_regressor.py b/src/tempor/methods/prediction/temporal/regression/plugin_seq2seq_regressor.py
index 93f4ad8a..7077f0c9 100644
--- a/src/tempor/methods/prediction/temporal/regression/plugin_seq2seq_regressor.py
+++ b/src/tempor/methods/prediction/temporal/regression/plugin_seq2seq_regressor.py
@@ -69,7 +69,7 @@ def _fit(
         self.model.fit(cl_dataset)
         return self
 
-    def _predict(  # type: ignore[override]
+    def _predict(
         self,
         data: dataset.PredictiveDataset,
         n_future_steps: int,
diff --git a/src/tempor/methods/time_to_event/__init__.py b/src/tempor/methods/time_to_event/__init__.py
index cc8f32b8..9b6329d9 100644
--- a/src/tempor/methods/time_to_event/__init__.py
+++ b/src/tempor/methods/time_to_event/__init__.py
@@ -41,7 +41,7 @@ def fit(self, data: dataset.BaseDataset, *args: Any, **kwargs: Any) -> Self:  #
     # specified here for unknown reasons. For now, we just ignore the type checking for these arguments with
     # `# type: ignore [no-untyped-def]`.
     @pydantic_utils.validate_arguments(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
-    def predict(  # type: ignore [no-untyped-def, override] # pylint: disable=arguments-differ
+    def predict(  # type: ignore [no-untyped-def] # pylint: disable=arguments-differ
         self,
         data: dataset.PredictiveDataset,
         horizons: data_typing.TimeIndex,
@@ -70,7 +70,7 @@ def predict_proba(self, data: dataset.PredictiveDataset, *args: Any, **kwargs: A
         )
 
     @abc.abstractmethod
-    def _predict(  # type: ignore[override]  # pylint: disable=arguments-differ
+    def _predict(  # pylint: disable=arguments-differ
         self, data: dataset.PredictiveDataset, horizons: data_typing.TimeIndex, *args: Any, **kwargs: Any
     ) -> samples.TimeSeriesSamples:  # pragma: no cover
         ...
diff --git a/src/tempor/methods/time_to_event/plugin_ddh.py b/src/tempor/methods/time_to_event/plugin_ddh.py
index 55a023a5..76f55146 100644
--- a/src/tempor/methods/time_to_event/plugin_ddh.py
+++ b/src/tempor/methods/time_to_event/plugin_ddh.py
@@ -102,7 +102,7 @@ def _fit(
         self.model.fit(processed_data, event_times, event_values)
         return self
 
-    def _predict(  # type: ignore[override]
+    def _predict(
         self,
         data: dataset.PredictiveDataset,
         horizons: data_typing.TimeIndex,
diff --git a/src/tempor/methods/time_to_event/plugin_ts_coxph.py b/src/tempor/methods/time_to_event/plugin_ts_coxph.py
index 878794df..b425a7fe 100644
--- a/src/tempor/methods/time_to_event/plugin_ts_coxph.py
+++ b/src/tempor/methods/time_to_event/plugin_ts_coxph.py
@@ -217,7 +217,7 @@ def _fit(
         self.model.fit(data, *args, **kwargs)
         return self
 
-    def _predict(  # type: ignore[override]
+    def _predict(
         self,
         data: dataset.PredictiveDataset,
         horizons: data_typing.TimeIndex,
diff --git a/src/tempor/methods/time_to_event/plugin_ts_xgb.py b/src/tempor/methods/time_to_event/plugin_ts_xgb.py
index 2d67d7c9..61099c49 100644
--- a/src/tempor/methods/time_to_event/plugin_ts_xgb.py
+++ b/src/tempor/methods/time_to_event/plugin_ts_xgb.py
@@ -381,7 +381,7 @@ def _fit(
         self.model.fit(data, *args, **kwargs)
         return self
 
-    def _predict(  # type: ignore[override]
+    def _predict(
         self,
         data: dataset.PredictiveDataset,
         horizons: data_typing.TimeIndex,
diff --git a/src/tempor/methods/treatments/temporal/classification/plugin_crn_classifier.py b/src/tempor/methods/treatments/temporal/classification/plugin_crn_classifier.py
index a4f7b28b..ce8d63cd 100644
--- a/src/tempor/methods/treatments/temporal/classification/plugin_crn_classifier.py
+++ b/src/tempor/methods/treatments/temporal/classification/plugin_crn_classifier.py
@@ -63,7 +63,7 @@ def _fit(
         self.model.fit(cl_dataset)
         return self
 
-    def _predict(  # type: ignore[override]  # pylint: disable=arguments-differ
+    def _predict(  # pylint: disable=arguments-differ
         self,
         data: dataset.PredictiveDataset,
         horizons: List[List[float]],
@@ -86,7 +86,7 @@ def _predict(  # type: ignore[override]  # pylint: disable=arguments-differ
         preds = _from_clv2_time_series(preds_cl.to_multi_index_dataframe())
         return samples.TimeSeriesSamples.from_dataframe(preds)
 
-    def _predict_counterfactuals(  # type: ignore[override]  # pylint: disable=arguments-differ
+    def _predict_counterfactuals(  # pylint: disable=arguments-differ
         self,
         data: dataset.PredictiveDataset,
         horizons: List[List[float]],
diff --git a/src/tempor/methods/treatments/temporal/regression/plugin_crn_regressor.py b/src/tempor/methods/treatments/temporal/regression/plugin_crn_regressor.py
index 24eabf34..fabba0fb 100644
--- a/src/tempor/methods/treatments/temporal/regression/plugin_crn_regressor.py
+++ b/src/tempor/methods/treatments/temporal/regression/plugin_crn_regressor.py
@@ -67,7 +67,7 @@ def _fit(
         self.model.fit(cl_dataset)
         return self
 
-    def _predict(  # type: ignore[override]  # pylint: disable=arguments-differ
+    def _predict(  # pylint: disable=arguments-differ
         self,
         data: dataset.PredictiveDataset,
         horizons: List[List[float]],
@@ -90,7 +90,7 @@ def _predict(  # type: ignore[override]  # pylint: disable=arguments-differ
         preds = _from_clv2_time_series(preds_cl.to_multi_index_dataframe())
         return samples.TimeSeriesSamples.from_dataframe(preds)
 
-    def _predict_counterfactuals(  # type: ignore[override]  # pylint: disable=arguments-differ
+    def _predict_counterfactuals(  # pylint: disable=arguments-differ
         self,
         data: dataset.PredictiveDataset,
         horizons: List[List[float]],
diff --git a/src/tempor/metrics/__init__.py b/src/tempor/metrics/__init__.py
new file mode 100644
index 00000000..121c88fb
--- /dev/null
+++ b/src/tempor/metrics/__init__.py
@@ -0,0 +1,9 @@
+"""Package directory for TemporAI metrics."""
+
+from . import prediction, time_to_event, treatments
+
+__all__ = [
+    "prediction",
+    "time_to_event",
+    "treatments",
+]
diff --git a/src/tempor/metrics/metric.py b/src/tempor/metrics/metric.py
new file mode 100644
index 00000000..57fd9e7d
--- /dev/null
+++ b/src/tempor/metrics/metric.py
@@ -0,0 +1,214 @@
+"""Module containing the base class for metrics."""
+
+# pylint: disable=unnecessary-ellipsis
+
+import abc
+from typing import Any, Generator, List
+
+import numpy as np
+import rich.pretty
+
+import tempor.core.utils
+from tempor.core import plugins
+from tempor.data import data_typing
+
+from . import metric_typing
+
+
+class Metric(plugins.Plugin, abc.ABC):
+    """Metric abstract base class, defines the required methods."""
+
+    @property
+    @abc.abstractmethod
+    def direction(self) -> metric_typing.MetricDirection:  # pragma: no cover
+        """The direction of the metric"""
+        ...
+
+    def evaluate(self, actual: Any, predicted: Any, *args: Any, **kwargs: Any) -> Any:  # pragma: no cover
+        """The metric evaluation call.
+
+        Args:
+            actual (Any): Actual value(s).
+            predicted (Any): Predicted value(s).
+            *args (Any): Additional positional arguments.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            Any: Evaluated metric value(s)
+        """
+        self._validate(actual, predicted)
+        return self._evaluate(actual, predicted, *args, **kwargs)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """A convenience method to call `evaluate` directly."""
+        return self.evaluate(*args, **kwargs)
+
+    @abc.abstractmethod
+    def _evaluate(self, actual: Any, predicted: Any, *args: Any, **kwargs: Any) -> Any:  # pragma: no cover
+        """The metric evaluation call *implementation* to be overridden in derived classes.
+
+        Args:
+            actual (Any): Actual value(s).
+            predicted (Any): Predicted value(s).
+            *args (Any): Additional positional arguments.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            Any: Evaluated metric value(s).
+        """
+        ...
+
+    def _validate(self, actual: Any, predicted: Any) -> None:
+        """Input validation. Can be overridden in derived classes, but a call to `super()._validate(...)` should
+        be preserved.
+
+        Args:
+            actual (Any): Actual value(s).
+            predicted (Any): Predicted value(s).
+        """
+        if actual is None or predicted is None:
+            raise ValueError("The actual values and the predicted values must not be `None`.")
+
+    def __rich_repr__(self) -> Generator:
+        """A `rich` representation of the class. The ``description`` field is auto-generated from the class and init
+        docstrings.
+
+        Yields:
+            Generator: The fields and their values fed to `rich`.
+        """
+        yield "name", self.name
+        yield "description", tempor.core.utils.make_description_from_doc(self)
+
+    def __repr__(self) -> str:
+        """The `repr()` representation of the class.
+
+        Returns:
+            str: The representation.
+        """
+        return rich.pretty.pretty_repr(self)
+
+
+# TODO: Multi-feature cases.
+# TODO: Typing of arguments may change.
+# TODO: Update the abstract methods for each case properly.
+
+
+class OneOffPredictionMetric(Metric):
+    """Metric abstract base class for the one-off prediction task."""
+
+    # Overridden for type hinting.
+    @abc.abstractmethod
+    def _evaluate(
+        self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any
+    ) -> float:  # pragma: no cover # noqa: D102
+        ...
+
+
+# Overrides for type hinting and docstrings.
+class OneOffClassificationMetric(OneOffPredictionMetric):
+    """Metric abstract base class for the one-off prediction task, classification case."""
+
+    def evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        """The metric evaluation call.
+
+        ``actual`` and ``predicted`` are expected to be numpy arrays (sample in the 0th dimension).
+
+        ``predicted`` must be the probabilities in this case.
+
+        Args:
+            actual (np.ndarray): Actual value(s).
+            predicted (np.ndarray): Predicted value(s).
+            *args (Any): Additional positional arguments.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            float: Evaluated metric value.
+        """
+        return super().evaluate(actual, predicted, *args, **kwargs)
+
+    @abc.abstractmethod
+    def _evaluate(
+        self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any
+    ) -> float:  # pragma: no cover # noqa: D102
+        ...
+
+
+plugins.register_plugin_category("prediction.one_off.classification", OneOffClassificationMetric, plugin_type="metric")
+
+
+# Overrides for type hinting and docstrings.
+class OneOffRegressionMetric(OneOffPredictionMetric):
+    """Metric abstract base class for the one-off prediction task, regression case."""
+
+    def evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        """The metric evaluation call.
+
+        ``actual`` and ``predicted`` are expected to be numpy arrays (sample in the 0th dimension).
+
+        Args:
+            actual (np.ndarray): Actual value(s).
+            predicted (np.ndarray): Predicted value(s).
+            *args (Any): Additional positional arguments.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            float: Evaluated metric value.
+        """
+        return super().evaluate(actual, predicted, *args, **kwargs)
+
+    @abc.abstractmethod
+    def _evaluate(
+        self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any
+    ) -> float:  # pragma: no cover # noqa: D102
+        ...
+
+
+plugins.register_plugin_category("prediction.one_off.regression", OneOffRegressionMetric, plugin_type="metric")
+
+
+# Overrides for type hinting and docstrings.
+class TimeToEventMetric(Metric):
+    """Metric abstract base class for the time-to-event (survival) task."""
+
+    def evaluate(  # pylint: disable=arguments-differ
+        self,
+        actual: metric_typing.EventArrayTimeArray,
+        predicted: np.ndarray,
+        horizons: data_typing.TimeIndex,
+        *args: Any,
+        **kwargs: Any,
+    ) -> List[float]:
+        """The metric evaluation call.
+
+        Args:
+            actual (metric_typing.EventArrayTimeArray):
+                A tuple of two numpy arrays: the event values array and the event times array,
+                for the actual event vales.
+            predicted (np.ndarray):
+                A numpy array of shape ``(n_samples, n_horizons_timesteps, n_features)``
+                with the predicted risk estimates.
+            horizons (data_typing.TimeIndex):
+                List of horizons time points.
+            *args (Any):
+                Additional positional arguments.
+            **kwargs (Any):
+                Additional keyword arguments.
+
+        Returns:
+            List[float]: The metric values for each horizon time point.
+        """
+        return super().evaluate(actual, predicted, horizons, *args, **kwargs)
+
+    @abc.abstractmethod
+    def _evaluate(  # pylint: disable=arguments-differ
+        self,
+        actual: metric_typing.EventArrayTimeArray,
+        predicted: np.ndarray,
+        horizons: data_typing.TimeIndex,
+        *args: Any,
+        **kwargs: Any,
+    ) -> List[float]:  # pragma: no cover # noqa: D102
+        ...
+
+
+plugins.register_plugin_category("time_to_event", TimeToEventMetric, plugin_type="metric")
diff --git a/src/tempor/metrics/metric_typing.py b/src/tempor/metrics/metric_typing.py
new file mode 100644
index 00000000..83bbb2be
--- /dev/null
+++ b/src/tempor/metrics/metric_typing.py
@@ -0,0 +1,16 @@
+"""Typing related to metrics."""
+
+from typing import Tuple
+
+import numpy as np
+from typing_extensions import Literal
+
+MetricDirection = Literal["minimize", "maximize"]
+"""The direction of the metric that represents the optimization goal (the "good" direction):
+``"minimize"`` or "``maximize``".
+"""
+
+EventArrayTimeArray = Tuple[np.ndarray, np.ndarray]
+"""Type hint used in time-to-event metrics, a tuple of two numpy arrays:
+the event values array and the event times array.
+"""
diff --git a/src/tempor/metrics/prediction/__init__.py b/src/tempor/metrics/prediction/__init__.py
new file mode 100644
index 00000000..e347ba3a
--- /dev/null
+++ b/src/tempor/metrics/prediction/__init__.py
@@ -0,0 +1,8 @@
+"""Package directory for prediction metrics."""
+
+from . import one_off, temporal
+
+__all__ = [
+    "one_off",
+    "temporal",
+]
diff --git a/src/tempor/metrics/prediction/one_off/__init__.py b/src/tempor/metrics/prediction/one_off/__init__.py
new file mode 100644
index 00000000..b3bf0959
--- /dev/null
+++ b/src/tempor/metrics/prediction/one_off/__init__.py
@@ -0,0 +1,9 @@
+"""Metrics for the one-off prediction tasks."""
+
+from tempor.core import plugins
+
+plugins.importing.import_plugins(__file__)
+
+__all__ = [  # pyright: ignore
+    *plugins.importing.gather_modules_names(__file__),
+]
diff --git a/src/tempor/metrics/prediction/one_off/plugin_builtin_classification.py b/src/tempor/metrics/prediction/one_off/plugin_builtin_classification.py
new file mode 100644
index 00000000..44b84dfc
--- /dev/null
+++ b/src/tempor/metrics/prediction/one_off/plugin_builtin_classification.py
@@ -0,0 +1,419 @@
+"""Module with built-in metric plugins for the category: prediction -> one-off -> classification."""
+
+from typing import Any, List, Tuple, cast
+
+import numpy as np
+import sklearn.metrics
+from sklearn.preprocessing import label_binarize
+
+from tempor.core import plugins
+from tempor.metrics import metric, metric_typing
+
+
+def _cast_to_y_pred(y_pred_proba: np.ndarray) -> np.ndarray:
+    """Turn an array of class probabilities ``y_pred_proba`` into an array of class predictions ``y_pred``.
+
+    Args:
+        y_pred_proba (np.ndarray): Predicted probabilities.
+
+    Returns:
+        np.ndarray: Predicted classes.
+    """
+    if y_pred_proba.ndim == 2:
+        # Interpret as probabilities.
+        y_pred = np.argmax(np.asarray(y_pred_proba), axis=1)
+    elif y_pred_proba.ndim == 1:
+        # Interpret as class labels.
+        y_pred = np.asarray(y_pred_proba)
+    else:
+        raise ValueError(f"Invalid shape of y_pred_proba: {y_pred_proba.shape}")
+    return y_pred
+
+
+def _get_y_pred_proba_hlpr(y_pred_proba: np.ndarray, nclasses: int) -> np.ndarray:
+    """A helper utility for inferring the correct y_pred_proba for multiclass situations, specifically in the case
+    of binary classification. See source code for the specifics.
+
+    Args:
+        y_pred_proba (np.ndarray): Predicted probabilities.
+        nclasses (int): Number of classes.
+
+    Returns:
+        np.ndarray: The correctly inferred ``y_pred_proba``.
+    """
+    if nclasses == 2:
+        if len(y_pred_proba.shape) < 2:
+            return y_pred_proba
+
+        if y_pred_proba.shape[1] == 2:
+            return y_pred_proba[:, 1]
+
+    return y_pred_proba
+
+
+def _prep_auc_multiclass(y_test: np.ndarray, y_pred_proba: np.ndarray) -> Tuple[np.ndarray, np.ndarray, int, List]:
+    y_test = np.asarray(y_test)
+    y_pred_proba = np.asarray(y_pred_proba)
+
+    nnan = sum(np.ravel(np.isnan(y_pred_proba)))
+    if nnan:
+        raise ValueError("NaNs in predictions, aborting.")
+
+    n_classes = len(set(np.ravel(y_test)))
+    classes = sorted(set(np.ravel(y_test)))
+
+    y_pred_proba_tmp = _get_y_pred_proba_hlpr(y_pred_proba, n_classes)
+
+    return y_test, y_pred_proba_tmp, n_classes, classes
+
+
+def _evaluate_aucroc_multiclass(
+    y_test: np.ndarray,
+    y_pred_proba: np.ndarray,
+) -> float:
+    """Helper for evaluating AUCROC for any number of classes."""
+
+    y_test, y_pred_proba_tmp, n_classes, classes = _prep_auc_multiclass(y_test, y_pred_proba)
+
+    if n_classes > 2:
+        fpr = dict()
+        tpr = dict()
+        precision = dict()
+        recall = dict()
+        roc_auc: dict = dict()
+
+        y_test = cast(np.ndarray, label_binarize(y_test, classes=classes, sparse_output=False))
+
+        fpr["micro"], tpr["micro"], _ = sklearn.metrics.roc_curve(y_test.ravel(), y_pred_proba_tmp.ravel())
+        roc_auc["micro"] = sklearn.metrics.auc(fpr["micro"], tpr["micro"])
+        precision["micro"], recall["micro"], _ = sklearn.metrics.precision_recall_curve(
+            y_test.ravel(), y_pred_proba_tmp.ravel()
+        )
+
+        aucroc = roc_auc["micro"]
+    else:
+        aucroc = sklearn.metrics.roc_auc_score(np.ravel(y_test), y_pred_proba_tmp, multi_class="ovr")
+
+    return float(aucroc)
+
+
+def _evaluate_aucprc_multiclass(
+    y_test: np.ndarray,
+    y_pred_proba: np.ndarray,
+) -> float:
+    """Helper for evaluating AUCPRC for any number of classes."""
+
+    y_test, y_pred_proba_tmp, n_classes, classes = _prep_auc_multiclass(y_test, y_pred_proba)
+
+    if n_classes > 2:
+        fpr = dict()
+        tpr = dict()
+        precision = dict()
+        recall = dict()
+        average_precision = dict()
+
+        y_test = cast(np.ndarray, label_binarize(y_test, classes=classes, sparse_output=False))
+
+        fpr["micro"], tpr["micro"], _ = sklearn.metrics.roc_curve(y_test.ravel(), y_pred_proba_tmp.ravel())
+        precision["micro"], recall["micro"], _ = sklearn.metrics.precision_recall_curve(
+            y_test.ravel(), y_pred_proba_tmp.ravel()
+        )
+        average_precision["micro"] = sklearn.metrics.average_precision_score(y_test, y_pred_proba_tmp, average="micro")
+
+        aucprc = average_precision["micro"]
+    else:
+        aucprc = sklearn.metrics.average_precision_score(np.ravel(y_test), y_pred_proba_tmp)
+
+    return float(aucprc)
+
+
+@plugins.register_plugin(name="accuracy", category="prediction.one_off.classification", plugin_type="metric")
+class AccuracyOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Accuracy classification score."""
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.accuracy_score(actual, _cast_to_y_pred(predicted)),
+        )
+
+
+@plugins.register_plugin(name="f1_score_micro", category="prediction.one_off.classification", plugin_type="metric")
+class F1ScoreMicroOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """F1 score is a harmonic mean of the precision and recall. This version uses the ``"micro"`` average: calculate
+    metrics globally by counting the total true positives, false negatives and false positives.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.f1_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="micro",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="f1_score_macro", category="prediction.one_off.classification", plugin_type="metric")
+class F1ScoreMacroOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """F1 score is a harmonic mean of the precision and recall. This version uses the ``"macro"`` average: calculate
+    metrics for each label, and find their unweighted mean. This does not take label imbalance into account.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.f1_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="macro",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="f1_score_weighted", category="prediction.one_off.classification", plugin_type="metric")
+class F1ScoreWeightedOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """F1 score is a harmonic mean of the precision and recall. This version uses the ``"weighted"`` average: calculate
+    metrics for each label, and find their average weighted by support (the number of true instances for each label).
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.f1_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="weighted",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="kappa", category="prediction.one_off.classification", plugin_type="metric")
+class KappaOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Computes Cohen's kappa, a score that expresses the level of agreement between two annotators on a classification
+    problem.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.cohen_kappa_score(actual, _cast_to_y_pred(predicted)),
+        )
+
+
+@plugins.register_plugin(name="kappa_quadratic", category="prediction.one_off.classification", plugin_type="metric")
+class KappaQuadraticOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Computes Cohen's kappa, a score that expresses the level of agreement between two annotators on a classification
+    problem. Weighted using the `"quadratic"` weighting.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.cohen_kappa_score(actual, _cast_to_y_pred(predicted), weights="quadratic"),
+        )
+
+
+@plugins.register_plugin(name="recall_micro", category="prediction.one_off.classification", plugin_type="metric")
+class RecallMicroOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Recall is defined as the number of true positives over the number of true positives plus the number of false
+    negatives. This version (micro) calculates metrics globally by counting the total true positives.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.recall_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="micro",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="recall_macro", category="prediction.one_off.classification", plugin_type="metric")
+class RecallMacroOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Recall is defined as the number of true positives over the number of true positives plus the number of false
+    negatives. This version (macro) calculates metrics for each label, and finds their unweighted mean.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.recall_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="macro",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="recall_weighted", category="prediction.one_off.classification", plugin_type="metric")
+class RecallWeightedOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Recall is defined as the number of true positives over the number of true positives plus the number of false
+    negatives. This version(weighted) calculates metrics for each label, and find their average weighted by support.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.recall_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="weighted",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="precision_micro", category="prediction.one_off.classification", plugin_type="metric")
+class PrecisionMicroOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Precision is defined as the number of true positives over the number of true positives plus the number of false
+    positives. This version (micro) calculates metrics globally by counting the total true positives.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.precision_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="micro",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="precision_macro", category="prediction.one_off.classification", plugin_type="metric")
+class PrecisionMacroOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Precision is defined as the number of true positives over the number of true positives plus the number of
+    false positives. This version (macro) calculates metrics for each label, and finds their unweighted mean.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.precision_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="macro",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="precision_weighted", category="prediction.one_off.classification", plugin_type="metric")
+class PrecisionWeightedOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """Precision is defined as the number of true positives over the number of true positives plus the number of
+    false positives. This version (weighted) calculates metrics for each label, and find their average weighted
+    by support.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.precision_score(
+                actual,
+                _cast_to_y_pred(predicted),
+                average="weighted",
+                zero_division=0,
+            ),
+        )
+
+
+@plugins.register_plugin(name="mcc", category="prediction.one_off.classification", plugin_type="metric")
+class MccOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """The Matthews Correlation Coefficient is used in machine learning as a measure of the quality of binary and
+    multiclass classifications. It takes into account true and false positives and negatives and is generally
+    regarded as a balanced measure which can be used even if the classes are of very different sizes.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.matthews_corrcoef(actual, _cast_to_y_pred(predicted)),
+        )
+
+
+@plugins.register_plugin(name="aucprc", category="prediction.one_off.classification", plugin_type="metric")
+class AucPrcOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """The average precision summarizes a precision-recall curve as the weighted mean of precisions achieved at each
+    threshold, with the increase in recall from the previous threshold used as the weight.
+    """
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return _evaluate_aucprc_multiclass(actual, predicted)
+
+
+@plugins.register_plugin(name="aucroc", category="prediction.one_off.classification", plugin_type="metric")
+class AucRocOneOffClassificationMetric(metric.OneOffClassificationMetric):
+    """The Area Under the Receiver Operating Characteristic Curve (ROC AUC) from prediction scores."""
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return _evaluate_aucroc_multiclass(actual, predicted)
diff --git a/src/tempor/metrics/prediction/one_off/plugin_builtin_regression.py b/src/tempor/metrics/prediction/one_off/plugin_builtin_regression.py
new file mode 100644
index 00000000..1ff197da
--- /dev/null
+++ b/src/tempor/metrics/prediction/one_off/plugin_builtin_regression.py
@@ -0,0 +1,54 @@
+"""Module with built-in metric plugins for the category: prediction -> one-off -> regression."""
+
+from typing import Any, cast
+
+import numpy as np
+import sklearn.metrics
+
+from tempor.core import plugins
+from tempor.metrics import metric, metric_typing
+
+
+@plugins.register_plugin(name="mse", category="prediction.one_off.regression", plugin_type="metric")
+class MseOneOffRegressionMetric(metric.OneOffRegressionMetric):
+    """Mean squared error regression metric"""
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "minimize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.mean_squared_error(actual, predicted),
+        )
+
+
+@plugins.register_plugin(name="mae", category="prediction.one_off.regression", plugin_type="metric")
+class MaeOneOffRegressionMetric(metric.OneOffRegressionMetric):
+    """Mean absolute error regression metric"""
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "minimize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.mean_absolute_error(actual, predicted),
+        )
+
+
+@plugins.register_plugin(name="r2", category="prediction.one_off.regression", plugin_type="metric")
+class R2OneOffRegressionMetric(metric.OneOffRegressionMetric):
+    """R^2 (coefficient of determination) score regression metric"""
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:
+        return cast(
+            float,
+            sklearn.metrics.r2_score(actual, predicted),
+        )
diff --git a/src/tempor/metrics/prediction/temporal/__init__.py b/src/tempor/metrics/prediction/temporal/__init__.py
new file mode 100644
index 00000000..accc920e
--- /dev/null
+++ b/src/tempor/metrics/prediction/temporal/__init__.py
@@ -0,0 +1,3 @@
+"""Metrics for the temporal prediction tasks."""
+
+# TODO: This.
diff --git a/src/tempor/metrics/time_to_event/__init__.py b/src/tempor/metrics/time_to_event/__init__.py
new file mode 100644
index 00000000..6c9c93dc
--- /dev/null
+++ b/src/tempor/metrics/time_to_event/__init__.py
@@ -0,0 +1,9 @@
+"""Metrics for the time-to-event (survival) analysis tasks."""
+
+from tempor.core import plugins
+
+plugins.importing.import_plugins(__file__)
+
+__all__ = [  # pyright: ignore
+    *plugins.importing.gather_modules_names(__file__),
+]
diff --git a/src/tempor/benchmarks/metrics.py b/src/tempor/metrics/time_to_event/_metric_impl.py
similarity index 93%
rename from src/tempor/benchmarks/metrics.py
rename to src/tempor/metrics/time_to_event/_metric_impl.py
index 9f277be6..be13274d 100644
--- a/src/tempor/benchmarks/metrics.py
+++ b/src/tempor/metrics/time_to_event/_metric_impl.py
@@ -1,6 +1,6 @@
-"""Module with any additional metrics."""
+"""Module with implementations for time-to-event (survival) analysis metric implementations."""
 
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import sklearn
@@ -97,7 +97,7 @@ def _check_times(test_time: np.ndarray, times: np.ndarray) -> np.ndarray:
         dtype=test_time.dtype,
         **dict(input_name="times") if _has_input_name else dict(),  # pyright: ignore
     )
-    times = np.unique(times)
+    times = np.unique(times)  # pyright: ignore
 
     if times.max() >= test_time.max() or times.min() < test_time.min():
         raise ValueError(
@@ -736,3 +736,51 @@ def concordance_index_ipcw(
     w = np.square(ipcw)
 
     return _estimate_concordance_index(test_event, test_time, estimate_arr, w, tied_tol)
+
+
+# --- Directly usable functions for the metrics. ---
+
+
+def compute_c_index(
+    training_array_struct: np.ndarray, testing_array_struct: np.ndarray, predictions: np.ndarray, horizons: List[float]
+) -> List[float]:
+    """Compute the IPCW concordance index.
+
+    Args:
+        training_array_struct (np.ndarray): Training data as a structured array.
+        testing_array_struct (np.ndarray): Testing data as a structured array.
+        predictions (np.ndarray): Predictions.
+        horizons (List[float]): Evaluation horizons.
+
+    Returns:
+        List[float]: List of metric values for each horizon.
+    """
+    metrics: List[float] = []
+    for horizon_idx, horizon_time in enumerate(horizons):
+        predictions_at_horizon_time = predictions[:, horizon_idx, :].reshape((-1,))
+        out = concordance_index_ipcw(
+            training_array_struct, testing_array_struct, predictions_at_horizon_time, float(horizon_time)
+        )
+        metrics.append(out[0])
+    return metrics
+
+
+def compute_brier_score(
+    training_array_struct: np.ndarray, testing_array_struct: np.ndarray, predictions: np.ndarray, horizons: List[float]
+) -> List[float]:
+    """Compute the time-dependent Brier score.
+
+    Args:
+        training_array_struct (np.ndarray): Training data as a structured array.
+        testing_array_struct (np.ndarray): Testing data as a structured array.
+        predictions (np.ndarray): Predictions.
+        horizons (List[float]): Evaluation horizons.
+
+    Returns:
+        List[float]: List of metric values for each horizon.
+    """
+    predictions = predictions.reshape((predictions.shape[0], predictions.shape[1]))
+    times, scores = brier_score(  # pylint: disable=unused-variable
+        training_array_struct, testing_array_struct, predictions, horizons
+    )
+    return scores.tolist()
diff --git a/src/tempor/metrics/time_to_event/plugin_builtin_time_to_event.py b/src/tempor/metrics/time_to_event/plugin_builtin_time_to_event.py
new file mode 100644
index 00000000..bb5913b8
--- /dev/null
+++ b/src/tempor/metrics/time_to_event/plugin_builtin_time_to_event.py
@@ -0,0 +1,176 @@
+"""Module with built-in metric plugins for the time-to-event (survival) analysis task."""
+from typing import Any, Callable, List
+
+import numpy as np
+
+from tempor.core import plugins
+from tempor.data import data_typing
+from tempor.metrics import metric, metric_typing
+
+from . import _metric_impl
+
+
+def _prep_surv_metric_data(
+    metric_func: Callable,
+    y_test: np.ndarray,
+    t_test: np.ndarray,
+    predictions_array: np.ndarray,
+    horizons: List[float],
+    y_train: np.ndarray,
+    t_train: np.ndarray,
+) -> List[float]:
+    """Prepare the input arrays for the survival metrics and return the metric per horizon.
+
+    Args:
+        metric_func (Callable): Metric function to be used.
+        y_test (np.ndarray): Event values, test set.
+        t_test (np.ndarray): Event time points, test set.
+        predictions_array (np.ndarray): Predictions array of shape ``(n_samples, n_horizons_timesteps, n_features)``
+        horizons (List[float]): List of horizons time points.
+        y_train (np.ndarray): Event values, training set.
+        t_train (np.ndarray): Event time points, training set.
+
+    Returns:
+        List[float]: Metric values for each horizon time point.
+    """
+    y_train_struct = _metric_impl.create_structured_array(y_train, t_train)
+    y_test_struct = _metric_impl.create_structured_array(y_test, t_test)
+
+    metric_per_horizon: List[float] = metric_func(y_train_struct, y_test_struct, predictions_array, horizons)
+
+    return metric_per_horizon
+
+
+# TODO: support non-float horizons.
+
+
+@plugins.register_plugin(name="c_index", category="time_to_event", plugin_type="metric")
+class CIndexTimeToEventMetric(metric.TimeToEventMetric):
+    """IPCW concordance index metric for time-to-event (survival) analysis tasks."""
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "maximize"
+
+    # Override to update the extra argument `actual_train`.
+    def evaluate(  # pylint: disable=arguments-differ
+        self,
+        actual: metric_typing.EventArrayTimeArray,
+        predicted: np.ndarray,
+        horizons: data_typing.TimeIndex,
+        actual_train: metric_typing.EventArrayTimeArray,
+        *args: Any,
+        **kwargs: Any,
+    ) -> List[float]:
+        """The metric evaluation call.
+
+        Args:
+            actual (metric_typing.EventArrayTimeArray):
+                A tuple of two numpy arrays: the event values array and the event times array,
+                for the actual event vales.
+            predicted (np.ndarray):
+                A numpy array of shape ``(n_samples, n_horizons_timesteps, n_features)``
+                with the predicted risk estimates.
+            horizons (data_typing.TimeIndex):
+                List of horizons time points.
+            actual_train (metric_typing.EventArrayTimeArray):
+                A tuple of two numpy arrays: the event values array and the event times array,
+                for the actual event vales - in the training set.
+            *args (Any):
+                Additional positional arguments.
+            **kwargs (Any):
+                Additional keyword arguments.
+
+        Returns:
+            List[float]: The metric values for each horizon time point.
+        """
+        return super().evaluate(actual, predicted, horizons, actual_train, *args, **kwargs)
+
+    def _evaluate(  # pylint: disable=arguments-differ
+        self,
+        actual: metric_typing.EventArrayTimeArray,
+        predicted: np.ndarray,
+        horizons: data_typing.TimeIndex,
+        actual_train: metric_typing.EventArrayTimeArray,
+        *args: Any,
+        **kwargs: Any,
+    ) -> List[float]:
+        y_test, t_test = actual
+        y_train, t_train = actual_train
+
+        metric_per_horizon = _prep_surv_metric_data(
+            _metric_impl.compute_c_index,
+            y_test,
+            t_test,
+            predicted,
+            horizons,  # type: ignore
+            y_train,
+            t_train,
+        )
+        return metric_per_horizon
+
+
+@plugins.register_plugin(name="brier_score", category="time_to_event", plugin_type="metric")
+class BrierScoreTimeToEventMetric(metric.TimeToEventMetric):
+    """Time-dependent Brier score metric for time-to-event (survival) analysis tasks."""
+
+    @property
+    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102
+        return "minimize"
+
+    # Override to update the extra argument `actual_train`.
+    def evaluate(  # pylint: disable=arguments-differ
+        self,
+        actual: metric_typing.EventArrayTimeArray,
+        predicted: np.ndarray,
+        horizons: data_typing.TimeIndex,
+        actual_train: metric_typing.EventArrayTimeArray,
+        *args: Any,
+        **kwargs: Any,
+    ) -> List[float]:
+        """The metric evaluation call.
+
+        Args:
+            actual (metric_typing.EventArrayTimeArray):
+                A tuple of two numpy arrays: the event values array and the event times array,
+                for the actual event vales.
+            predicted (np.ndarray):
+                A numpy array of shape ``(n_samples, n_horizons_timesteps, n_features)``
+                with the predicted risk estimates.
+            horizons (data_typing.TimeIndex):
+                List of horizons time points.
+            actual_train (metric_typing.EventArrayTimeArray):
+                A tuple of two numpy arrays: the event values array and the event times array,
+                for the actual event vales - in the training set.
+            *args (Any):
+                Additional positional arguments.
+            **kwargs (Any):
+                Additional keyword arguments.
+
+        Returns:
+            List[float]: The metric values for each horizon time point.
+        """
+        return super().evaluate(actual, predicted, horizons, actual_train, *args, **kwargs)
+
+    def _evaluate(  # pylint: disable=arguments-differ
+        self,
+        actual: metric_typing.EventArrayTimeArray,
+        predicted: np.ndarray,
+        horizons: data_typing.TimeIndex,
+        actual_train: metric_typing.EventArrayTimeArray,
+        *args: Any,
+        **kwargs: Any,
+    ) -> List[float]:
+        y_test, t_test = actual
+        y_train, t_train = actual_train
+
+        metric_per_horizon = _prep_surv_metric_data(
+            _metric_impl.compute_brier_score,
+            y_test,
+            t_test,
+            predicted,
+            horizons,  # type: ignore
+            y_train,
+            t_train,
+        )
+        return metric_per_horizon
diff --git a/src/tempor/metrics/treatments/__init__.py b/src/tempor/metrics/treatments/__init__.py
new file mode 100644
index 00000000..ca1e1e00
--- /dev/null
+++ b/src/tempor/metrics/treatments/__init__.py
@@ -0,0 +1,8 @@
+"""Package directory for treatment effects metrics."""
+
+from . import one_off, temporal
+
+__all__ = [
+    "one_off",
+    "temporal",
+]
diff --git a/src/tempor/metrics/treatments/one_off/__init__.py b/src/tempor/metrics/treatments/one_off/__init__.py
new file mode 100644
index 00000000..1b828e6d
--- /dev/null
+++ b/src/tempor/metrics/treatments/one_off/__init__.py
@@ -0,0 +1,1 @@
+"""Metrics for the one-off treatment effects estimation tasks."""
diff --git a/src/tempor/metrics/treatments/temporal/__init__.py b/src/tempor/metrics/treatments/temporal/__init__.py
new file mode 100644
index 00000000..735eb662
--- /dev/null
+++ b/src/tempor/metrics/treatments/temporal/__init__.py
@@ -0,0 +1,1 @@
+"""Metrics for the temporal treatment effects estimation tasks."""
diff --git a/src/tempor/models/ddh.py b/src/tempor/models/ddh.py
index fce5bf6b..06b021bb 100644
--- a/src/tempor/models/ddh.py
+++ b/src/tempor/models/ddh.py
@@ -404,7 +404,7 @@ def negative_log_likelihood(
 
         # Censored loss
         loss += torch.sum(torch.log(nn.ReLU()(1 - censored_cif) + constants.EPS))
-        return -loss / len(outcomes)
+        return -loss / len(outcomes)  # type: ignore [return-value]
 
     def ranking_loss(
         self,
@@ -424,9 +424,11 @@ def ranking_loss(
                 # must have a lower risk for that cause
                 if torch.sum(t > ti) > 0:
                     # TODO: When data are sorted in time -> wan we make it even faster?
-                    loss += torch.mean(torch.exp((cif_k[t > ti][:, ti] - ci[ti])) / self.sigma)
+                    loss += torch.mean(  # type: ignore [call-overload]
+                        torch.exp((cif_k[t > ti][:, ti] - ci[ti])) / self.sigma
+                    )
 
-        return loss / len(cif)
+        return loss / len(cif)  # type: ignore [return-value]
 
     def longitudinal_loss(self, longitudinal_prediction: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
         """Penalize error in the longitudinal predictions. This function is used to compute the error made by the RNN.
@@ -471,7 +473,8 @@ def total_loss(
         cif = [torch.cumsum(ok, 1) for ok in outcomes]
 
         return (
-            (1 - self.alpha - self.beta) * self.longitudinal_loss(longitudinal_prediction, x)
+            (1 - self.alpha - self.beta)
+            * self.longitudinal_loss(longitudinal_prediction, x)  # type: ignore [return-value]
             + self.alpha * self.ranking_loss(cif, t, e)
             + self.beta * self.negative_log_likelihood(outcomes, cif, t, e)
         )
@@ -505,7 +508,7 @@ def __init__(
         # RNN model for longitudinal data
         self.embedding: nn.Module
         if self.rnn_type == "LSTM":
-            self.embedding = nn.LSTM(  # type: ignore [no-untyped-call]
+            self.embedding = nn.LSTM(
                 input_dim,
                 hidden_rnn,
                 layers_rnn,
@@ -513,7 +516,7 @@ def __init__(
                 batch_first=True,
             )
         elif self.rnn_type == "RNN":
-            self.embedding = nn.RNN(  # type: ignore [no-untyped-call]
+            self.embedding = nn.RNN(
                 input_dim,
                 hidden_rnn,
                 layers_rnn,
@@ -522,7 +525,7 @@ def __init__(
                 nonlinearity="relu",
             )
         elif self.rnn_type == "GRU":
-            self.embedding = nn.GRU(  # type: ignore [no-untyped-call]
+            self.embedding = nn.GRU(
                 input_dim,
                 hidden_rnn,
                 layers_rnn,
diff --git a/src/tempor/models/ts_model.py b/src/tempor/models/ts_model.py
index 926af6e1..ee36ed0c 100644
--- a/src/tempor/models/ts_model.py
+++ b/src/tempor/models/ts_model.py
@@ -594,7 +594,7 @@ def __init__(
         }
 
         if mode in ["RNN", "LSTM", "GRU"]:
-            self.temporal_layer = temporal_models[mode](**temporal_params)  # type: ignore [no-untyped-call]
+            self.temporal_layer = temporal_models[mode](**temporal_params)
         elif mode == "MLSTM_FCN":
             self.temporal_layer = MLSTM_FCN(
                 c_in=n_temporal_units_in,
diff --git a/src/tempor/models/ts_ode.py b/src/tempor/models/ts_ode.py
index 649fa07d..ff7cba52 100644
--- a/src/tempor/models/ts_ode.py
+++ b/src/tempor/models/ts_ode.py
@@ -131,7 +131,7 @@ def __init__(
     ):
         """Model (encoder and Laplace representation func). Encodes observed trajectory into latent vector."""
         super(ReverseGRUEncoder, self).__init__()
-        self.gru = nn.GRU(n_units_in, n_units_hidden, 2, batch_first=True)  # type: ignore [no-untyped-call]
+        self.gru = nn.GRU(n_units_in, n_units_hidden, 2, batch_first=True)
         self.linear_out = nn.Linear(n_units_hidden, n_units_latent).to(device)
         nn.init.xavier_uniform_(self.linear_out.weight)
 
diff --git a/src/tempor/models/utils.py b/src/tempor/models/utils.py
index 007682a6..45ab7b29 100644
--- a/src/tempor/models/utils.py
+++ b/src/tempor/models/utils.py
@@ -60,7 +60,7 @@ def enable_reproducibility(
     torch.cuda.manual_seed_all(random_seed)
     # If enabled, force deterministic algorithms:
     if torch_use_deterministic_algorithms:
-        torch.use_deterministic_algorithms(True)  # type: ignore [no-untyped-call]
+        torch.use_deterministic_algorithms(True)
         if warn_cuda_env_vars and torch.version.cuda not in ("", None):
             cuda_version = Version(torch.version.cuda)
             if (cuda_version.major, cuda_version.minor) == (10, 1) and os.environ.get(
diff --git a/tutorials/data/tutorial01_data_format.ipynb b/tutorials/data/tutorial01_data_format.ipynb
index 17d5536d..20022b6b 100644
--- a/tutorials/data/tutorial01_data_format.ipynb
+++ b/tutorials/data/tutorial01_data_format.ipynb
@@ -53,7 +53,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -65,7 +65,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -180,7 +180,7 @@
        "sample_2   9               31       3.1        c"
       ]
      },
-     "execution_count": 2,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -230,7 +230,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -300,7 +300,7 @@
        "sample_2       300      -1.3"
       ]
      },
-     "execution_count": 3,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -346,7 +346,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -416,7 +416,7 @@
        "sample_2   (13, True)   (11, True)"
       ]
      },
-     "execution_count": 4,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -467,7 +467,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -480,7 +480,7 @@
        ")"
       ]
      },
-     "execution_count": 5,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -501,7 +501,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -601,7 +601,7 @@
        "sample_2   9               31       3.1        c"
       ]
      },
-     "execution_count": 6,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -612,7 +612,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -674,7 +674,7 @@
        "sample_2         300      -1.3"
       ]
      },
-     "execution_count": 7,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -685,7 +685,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -747,7 +747,7 @@
        "sample_2     (13, True)   (11, True)"
       ]
      },
-     "execution_count": 8,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -776,7 +776,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -785,7 +785,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -884,7 +884,7 @@
        "sample_2   9               31       3.1        c"
       ]
      },
-     "execution_count": 10,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -897,7 +897,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -919,7 +919,7 @@
        "        [-999.0, -999.0, -999.0]]], dtype=object)"
       ]
      },
-     "execution_count": 11,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -932,7 +932,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -953,7 +953,7 @@
        " sample_2   9               31       3.1        c]"
       ]
      },
-     "execution_count": 12,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -966,7 +966,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -987,7 +987,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1014,7 +1014,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1027,7 +1027,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1112,7 +1112,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1237,8 +1237,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/data/tutorial02_datasets.ipynb b/tutorials/data/tutorial02_datasets.ipynb
index 91cf0697..49914903 100644
--- a/tutorials/data/tutorial02_datasets.ipynb
+++ b/tutorials/data/tutorial02_datasets.ipynb
@@ -39,7 +39,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -87,7 +87,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -186,7 +186,7 @@
        "sample_2   9               31       3.1       111"
       ]
      },
-     "execution_count": 2,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -197,7 +197,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -256,7 +256,7 @@
        "sample_2       300      -1.3         0"
       ]
      },
-     "execution_count": 3,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -267,7 +267,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -322,7 +322,7 @@
        "sample_2   (13, True)   (11, True)"
       ]
      },
-     "execution_count": 4,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -346,7 +346,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -355,7 +355,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -367,7 +367,7 @@
        ")"
       ]
      },
-     "execution_count": 6,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -384,7 +384,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -484,7 +484,7 @@
        "sample_2   9               31       3.1       111"
       ]
      },
-     "execution_count": 7,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -495,7 +495,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -562,7 +562,7 @@
        "sample_2         300      -1.3         0"
       ]
      },
-     "execution_count": 8,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -587,7 +587,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -600,7 +600,7 @@
        ")"
       ]
      },
-     "execution_count": 9,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -618,7 +618,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -718,7 +718,7 @@
        "sample_2   9               31       3.1       111"
       ]
      },
-     "execution_count": 10,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -729,7 +729,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -791,7 +791,7 @@
        "sample_2         300      -1.3"
       ]
      },
-     "execution_count": 11,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -802,7 +802,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -859,7 +859,7 @@
        "sample_2           0"
       ]
      },
-     "execution_count": 12,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -884,7 +884,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -899,7 +899,7 @@
        ")"
       ]
      },
-     "execution_count": 13,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -917,7 +917,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1008,7 +1008,7 @@
        "sample_2   9               31       3.1"
       ]
      },
-     "execution_count": 14,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1019,7 +1019,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1086,7 +1086,7 @@
        "sample_2         300      -1.3         0"
       ]
      },
-     "execution_count": 15,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1097,7 +1097,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1179,7 +1179,7 @@
        "sample_2   9              111"
       ]
      },
-     "execution_count": 16,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1204,7 +1204,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1217,7 +1217,7 @@
        ")"
       ]
      },
-     "execution_count": 17,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1235,7 +1235,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1335,7 +1335,7 @@
        "sample_2   9               31       3.1       111"
       ]
      },
-     "execution_count": 18,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1346,7 +1346,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1413,7 +1413,7 @@
        "sample_2         300      -1.3         0"
       ]
      },
-     "execution_count": 19,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1424,7 +1424,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1486,7 +1486,7 @@
        "sample_2     (13, True)   (11, True)"
       ]
      },
-     "execution_count": 20,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1513,7 +1513,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1529,7 +1529,7 @@
        ")"
       ]
      },
-     "execution_count": 21,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1548,7 +1548,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1639,7 +1639,7 @@
        "sample_2   9               31       3.1"
       ]
      },
-     "execution_count": 22,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1650,7 +1650,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1717,7 +1717,7 @@
        "sample_2         300      -1.3         0"
       ]
      },
-     "execution_count": 23,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1728,7 +1728,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1810,7 +1810,7 @@
        "sample_2   9              111"
       ]
      },
-     "execution_count": 24,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1821,7 +1821,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1878,7 +1878,7 @@
        "sample_2     (13, True)"
       ]
      },
-     "execution_count": 25,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1905,7 +1905,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1921,7 +1921,7 @@
        ")"
       ]
      },
-     "execution_count": 26,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1940,7 +1940,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -2022,7 +2022,7 @@
        "sample_2   9               31"
       ]
      },
-     "execution_count": 27,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2033,7 +2033,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -2100,7 +2100,7 @@
        "sample_2         300      -1.3         0"
       ]
      },
-     "execution_count": 28,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2111,7 +2111,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -2193,7 +2193,7 @@
        "sample_2   9              3.1"
       ]
      },
-     "execution_count": 29,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2204,7 +2204,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -2286,7 +2286,7 @@
        "sample_2   9              111"
       ]
      },
-     "execution_count": 30,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2336,8 +2336,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/data/tutorial03_datasources.ipynb b/tutorials/data/tutorial03_datasources.ipynb
index acd4f3c1..ed9b322b 100644
--- a/tutorials/data/tutorial03_datasources.ipynb
+++ b/tutorials/data/tutorial03_datasources.ipynb
@@ -48,7 +48,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -69,16 +69,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<tempor.datasources.prediction.one_off.plugin_sine.SineDataSource at 0x7f18cd1f8670>"
+       "<tempor.datasources.prediction.one_off.plugin_sine.SineDataSource at 0x7ff8f84f6070>"
       ]
      },
-     "execution_count": 27,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -97,7 +97,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -117,7 +117,7 @@
        ")"
       ]
      },
-     "execution_count": 28,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -133,7 +133,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -290,7 +290,7 @@
        "[400 rows x 5 columns]"
       ]
      },
-     "execution_count": 29,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -301,7 +301,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -440,7 +440,7 @@
        "[80 rows x 4 columns]"
       ]
      },
-     "execution_count": 30,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -451,7 +451,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -551,7 +551,7 @@
        "[80 rows x 1 columns]"
       ]
      },
-     "execution_count": 31,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -569,16 +569,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<tempor.datasources.prediction.one_off.plugin_sine.SineDataSource at 0x7f18cd228cd0>"
+       "<tempor.datasources.prediction.one_off.plugin_sine.SineDataSource at 0x7ff8f8552d90>"
       ]
      },
-     "execution_count": 32,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -603,7 +603,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -646,7 +646,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 34,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -657,7 +657,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "SineDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "OneOffPredictionDataset(\n",
       "    time_series=TimeSeriesSamples([100, *, 5]),\n",
       "    static=StaticSamples([100, 4]),\n",
@@ -1086,7 +1092,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "GoogleStocksDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "OneOffPredictionDataset(\n",
       "    time_series=TimeSeriesSamples([50, *, 5]),\n",
       "    predictive=OneOffPredictionTaskData(targets=StaticSamples([50, 1]))\n",
@@ -1600,24 +1612,24 @@
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th></th>\n",
-       "      <th>post-supper_blood_glucose_measurement</th>\n",
-       "      <th>less-than-usual_meal_ingestion</th>\n",
-       "      <th>regular_insulin_dose</th>\n",
-       "      <th>more-than-usual_exercise_activity</th>\n",
+       "      <th>post-lunch_blood_glucose_measurement</th>\n",
+       "      <th>more-than-usual_meal_ingestion</th>\n",
        "      <th>unspecified_special_event</th>\n",
        "      <th>typical_exercise_activity</th>\n",
-       "      <th>nph_insulin_dose</th>\n",
        "      <th>less-than-usual_exercise_activity</th>\n",
-       "      <th>typical_meal_ingestion</th>\n",
-       "      <th>pre-breakfast_blood_glucose_measurement</th>\n",
-       "      <th>pre-lunch_blood_glucose_measurement</th>\n",
-       "      <th>pre-supper_blood_glucose_measurement</th>\n",
+       "      <th>post-supper_blood_glucose_measurement</th>\n",
+       "      <th>more-than-usual_exercise_activity</th>\n",
        "      <th>unspecified_blood_glucose_measurement</th>\n",
-       "      <th>post-lunch_blood_glucose_measurement</th>\n",
-       "      <th>more-than-usual_meal_ingestion</th>\n",
-       "      <th>post-breakfast_blood_glucose_measurement</th>\n",
-       "      <th>ultralente_insulin_dose</th>\n",
        "      <th>pre-snack_blood_glucose_measurement</th>\n",
+       "      <th>ultralente_insulin_dose</th>\n",
+       "      <th>less-than-usual_meal_ingestion</th>\n",
+       "      <th>pre-supper_blood_glucose_measurement</th>\n",
+       "      <th>regular_insulin_dose</th>\n",
+       "      <th>pre-breakfast_blood_glucose_measurement</th>\n",
+       "      <th>nph_insulin_dose</th>\n",
+       "      <th>post-breakfast_blood_glucose_measurement</th>\n",
+       "      <th>pre-lunch_blood_glucose_measurement</th>\n",
+       "      <th>typical_meal_ingestion</th>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>sample_idx</th>\n",
@@ -1648,19 +1660,19 @@
        "      <th>0</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>9.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>13.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>100.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>9.0</td>\n",
+       "      <td>100.0</td>\n",
+       "      <td>13.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1669,7 +1681,6 @@
        "      <th>1</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1678,8 +1689,9 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>119.0</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>119.0</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1695,12 +1707,12 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>123.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>123.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1711,19 +1723,19 @@
        "      <th>3</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>10.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>13.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>216.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>10.0</td>\n",
+       "      <td>216.0</td>\n",
+       "      <td>13.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1732,7 +1744,6 @@
        "      <th>4</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>2.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1743,6 +1754,7 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>2.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1781,12 +1793,12 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>145.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>145.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1797,19 +1809,19 @@
        "      <th>147</th>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>1.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>259.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>259.0</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1822,7 +1834,6 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1831,6 +1842,7 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1843,7 +1855,6 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1852,6 +1863,7 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1864,7 +1876,6 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
-       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1873,6 +1884,7 @@
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
+       "      <td>7.0</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
        "      <td>NaN</td>\n",
@@ -1884,19 +1896,229 @@
       ],
       "text/plain": [
        "TimeSeriesSamples with data:\n",
-       "                     post-supper_blood_glucose_measurement  ...  pre-snack_blood_glucose_measurement\n",
-       "sample_idx time_idx                                         ...                                     \n",
-       "0          0                                           NaN  ...                                  NaN\n",
-       "           1                                           NaN  ...                                  NaN\n",
-       "           2                                           NaN  ...                                  NaN\n",
-       "           3                                           NaN  ...                                  NaN\n",
-       "           4                                           NaN  ...                                  NaN\n",
-       "...                                                    ...  ...                                  ...\n",
-       "69         146                                         NaN  ...                                  NaN\n",
-       "           147                                         NaN  ...                                  NaN\n",
-       "           148                                         NaN  ...                                  NaN\n",
-       "           149                                         NaN  ...                                  NaN\n",
-       "           150                                         NaN  ...                                  NaN\n",
+       "                     post-lunch_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                         \n",
+       "0          0                                          NaN   \n",
+       "           1                                          NaN   \n",
+       "           2                                          NaN   \n",
+       "           3                                          NaN   \n",
+       "           4                                          NaN   \n",
+       "...                                                   ...   \n",
+       "69         146                                        NaN   \n",
+       "           147                                        NaN   \n",
+       "           148                                        NaN   \n",
+       "           149                                        NaN   \n",
+       "           150                                        NaN   \n",
+       "\n",
+       "                     more-than-usual_meal_ingestion  \\\n",
+       "sample_idx time_idx                                   \n",
+       "0          0                                    NaN   \n",
+       "           1                                    NaN   \n",
+       "           2                                    NaN   \n",
+       "           3                                    NaN   \n",
+       "           4                                    NaN   \n",
+       "...                                             ...   \n",
+       "69         146                                  NaN   \n",
+       "           147                                  NaN   \n",
+       "           148                                  NaN   \n",
+       "           149                                  NaN   \n",
+       "           150                                  NaN   \n",
+       "\n",
+       "                     unspecified_special_event  typical_exercise_activity  \\\n",
+       "sample_idx time_idx                                                         \n",
+       "0          0                               NaN                        NaN   \n",
+       "           1                               NaN                        NaN   \n",
+       "           2                               NaN                        NaN   \n",
+       "           3                               NaN                        NaN   \n",
+       "           4                               NaN                        NaN   \n",
+       "...                                        ...                        ...   \n",
+       "69         146                             NaN                        NaN   \n",
+       "           147                             NaN                        NaN   \n",
+       "           148                             NaN                        NaN   \n",
+       "           149                             NaN                        NaN   \n",
+       "           150                             NaN                        NaN   \n",
+       "\n",
+       "                     less-than-usual_exercise_activity  \\\n",
+       "sample_idx time_idx                                      \n",
+       "0          0                                       NaN   \n",
+       "           1                                       NaN   \n",
+       "           2                                       NaN   \n",
+       "           3                                       NaN   \n",
+       "           4                                       NaN   \n",
+       "...                                                ...   \n",
+       "69         146                                     NaN   \n",
+       "           147                                     NaN   \n",
+       "           148                                     NaN   \n",
+       "           149                                     NaN   \n",
+       "           150                                     NaN   \n",
+       "\n",
+       "                     post-supper_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                          \n",
+       "0          0                                           NaN   \n",
+       "           1                                           NaN   \n",
+       "           2                                           NaN   \n",
+       "           3                                           NaN   \n",
+       "           4                                           NaN   \n",
+       "...                                                    ...   \n",
+       "69         146                                         NaN   \n",
+       "           147                                         NaN   \n",
+       "           148                                         NaN   \n",
+       "           149                                         NaN   \n",
+       "           150                                         NaN   \n",
+       "\n",
+       "                     more-than-usual_exercise_activity  \\\n",
+       "sample_idx time_idx                                      \n",
+       "0          0                                       NaN   \n",
+       "           1                                       NaN   \n",
+       "           2                                       NaN   \n",
+       "           3                                       NaN   \n",
+       "           4                                       NaN   \n",
+       "...                                                ...   \n",
+       "69         146                                     NaN   \n",
+       "           147                                     NaN   \n",
+       "           148                                     NaN   \n",
+       "           149                                     NaN   \n",
+       "           150                                     NaN   \n",
+       "\n",
+       "                     unspecified_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                          \n",
+       "0          0                                           NaN   \n",
+       "           1                                           NaN   \n",
+       "           2                                         123.0   \n",
+       "           3                                           NaN   \n",
+       "           4                                           NaN   \n",
+       "...                                                    ...   \n",
+       "69         146                                       145.0   \n",
+       "           147                                         NaN   \n",
+       "           148                                         NaN   \n",
+       "           149                                         NaN   \n",
+       "           150                                         NaN   \n",
+       "\n",
+       "                     pre-snack_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                        \n",
+       "0          0                                         NaN   \n",
+       "           1                                         NaN   \n",
+       "           2                                         NaN   \n",
+       "           3                                         NaN   \n",
+       "           4                                         NaN   \n",
+       "...                                                  ...   \n",
+       "69         146                                       NaN   \n",
+       "           147                                       NaN   \n",
+       "           148                                       NaN   \n",
+       "           149                                       NaN   \n",
+       "           150                                       NaN   \n",
+       "\n",
+       "                     ultralente_insulin_dose  less-than-usual_meal_ingestion  \\\n",
+       "sample_idx time_idx                                                            \n",
+       "0          0                             NaN                             NaN   \n",
+       "           1                             NaN                             NaN   \n",
+       "           2                             NaN                             NaN   \n",
+       "           3                             NaN                             NaN   \n",
+       "           4                             NaN                             NaN   \n",
+       "...                                      ...                             ...   \n",
+       "69         146                           NaN                             NaN   \n",
+       "           147                           NaN                             NaN   \n",
+       "           148                           NaN                             NaN   \n",
+       "           149                           NaN                             NaN   \n",
+       "           150                           NaN                             NaN   \n",
+       "\n",
+       "                     pre-supper_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                         \n",
+       "0          0                                          NaN   \n",
+       "           1                                        119.0   \n",
+       "           2                                          NaN   \n",
+       "           3                                          NaN   \n",
+       "           4                                          NaN   \n",
+       "...                                                   ...   \n",
+       "69         146                                        NaN   \n",
+       "           147                                        NaN   \n",
+       "           148                                        NaN   \n",
+       "           149                                        NaN   \n",
+       "           150                                        NaN   \n",
+       "\n",
+       "                     regular_insulin_dose  \\\n",
+       "sample_idx time_idx                         \n",
+       "0          0                          9.0   \n",
+       "           1                          7.0   \n",
+       "           2                          NaN   \n",
+       "           3                         10.0   \n",
+       "           4                          2.0   \n",
+       "...                                   ...   \n",
+       "69         146                        NaN   \n",
+       "           147                        1.0   \n",
+       "           148                        NaN   \n",
+       "           149                        NaN   \n",
+       "           150                        NaN   \n",
+       "\n",
+       "                     pre-breakfast_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                            \n",
+       "0          0                                           100.0   \n",
+       "           1                                             NaN   \n",
+       "           2                                             NaN   \n",
+       "           3                                           216.0   \n",
+       "           4                                             NaN   \n",
+       "...                                                      ...   \n",
+       "69         146                                           NaN   \n",
+       "           147                                         259.0   \n",
+       "           148                                           NaN   \n",
+       "           149                                           NaN   \n",
+       "           150                                           NaN   \n",
+       "\n",
+       "                     nph_insulin_dose  \\\n",
+       "sample_idx time_idx                     \n",
+       "0          0                     13.0   \n",
+       "           1                      NaN   \n",
+       "           2                      NaN   \n",
+       "           3                     13.0   \n",
+       "           4                      NaN   \n",
+       "...                               ...   \n",
+       "69         146                    NaN   \n",
+       "           147                    7.0   \n",
+       "           148                    7.0   \n",
+       "           149                    7.0   \n",
+       "           150                    7.0   \n",
+       "\n",
+       "                     post-breakfast_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                             \n",
+       "0          0                                              NaN   \n",
+       "           1                                              NaN   \n",
+       "           2                                              NaN   \n",
+       "           3                                              NaN   \n",
+       "           4                                              NaN   \n",
+       "...                                                       ...   \n",
+       "69         146                                            NaN   \n",
+       "           147                                            NaN   \n",
+       "           148                                            NaN   \n",
+       "           149                                            NaN   \n",
+       "           150                                            NaN   \n",
+       "\n",
+       "                     pre-lunch_blood_glucose_measurement  \\\n",
+       "sample_idx time_idx                                        \n",
+       "0          0                                         NaN   \n",
+       "           1                                         NaN   \n",
+       "           2                                         NaN   \n",
+       "           3                                         NaN   \n",
+       "           4                                         NaN   \n",
+       "...                                                  ...   \n",
+       "69         146                                       NaN   \n",
+       "           147                                       NaN   \n",
+       "           148                                       NaN   \n",
+       "           149                                       NaN   \n",
+       "           150                                       NaN   \n",
+       "\n",
+       "                     typical_meal_ingestion  \n",
+       "sample_idx time_idx                          \n",
+       "0          0                            NaN  \n",
+       "           1                            NaN  \n",
+       "           2                            NaN  \n",
+       "           3                            NaN  \n",
+       "           4                            NaN  \n",
+       "...                                     ...  \n",
+       "69         146                          NaN  \n",
+       "           147                          NaN  \n",
+       "           148                          NaN  \n",
+       "           149                          NaN  \n",
+       "           150                          NaN  \n",
        "\n",
        "[18199 rows x 18 columns]"
       ]
@@ -2024,7 +2246,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "DummyTemporalPredictionDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "TemporalPredictionDataset(\n",
       "    time_series=TimeSeriesSamples([100, *, 5]),\n",
       "    static=StaticSamples([100, 3]),\n",
@@ -2460,7 +2688,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "PBCDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "TimeToEventAnalysisDataset(\n",
       "    time_series=TimeSeriesSamples([312, *, 14]),\n",
       "    static=StaticSamples([312, 1]),\n",
@@ -2726,19 +2960,47 @@
       ],
       "text/plain": [
        "TimeSeriesSamples with data:\n",
-       "                     drug  ascites  ...  prothrombin       age\n",
-       "sample_idx time_idx                 ...                       \n",
-       "1          0.569489   0.0      1.0  ...     0.136768  0.248058\n",
-       "           1.095170   0.0      1.0  ...     0.813132  0.248058\n",
-       "2          5.319790   0.0      1.0  ...     0.339677  1.292856\n",
-       "           6.261636   0.0      1.0  ...     0.339677  1.292856\n",
-       "           7.266455   0.0      1.0  ...     0.339677  1.292856\n",
-       "...                   ...      ...  ...          ...       ...\n",
-       "312        1.045888   1.0      0.0  ...     1.624769 -1.962482\n",
-       "           1.867265   1.0      0.0  ...     1.354223 -1.962482\n",
-       "           2.921367   1.0      0.0  ...     0.474950 -1.962482\n",
-       "           3.425145   1.0      0.0  ...    -0.066141 -1.962482\n",
-       "           3.989158   1.0      0.0  ...    -0.133778 -1.962482\n",
+       "                     drug  ascites  hepatomegaly  spiders  edema  histologic  \\\n",
+       "sample_idx time_idx                                                            \n",
+       "1          0.569489   0.0      1.0           1.0      1.0    1.0         3.0   \n",
+       "           1.095170   0.0      1.0           1.0      1.0    1.0         3.0   \n",
+       "2          5.319790   0.0      1.0           1.0      1.0    1.0         2.0   \n",
+       "           6.261636   0.0      1.0           1.0      1.0    1.0         2.0   \n",
+       "           7.266455   0.0      1.0           1.0      1.0    1.0         2.0   \n",
+       "...                   ...      ...           ...      ...    ...         ...   \n",
+       "312        1.045888   1.0      0.0           0.0      1.0    2.0         2.0   \n",
+       "           1.867265   1.0      0.0           0.0      1.0    2.0         1.0   \n",
+       "           2.921367   1.0      0.0           0.0      0.0    0.0         1.0   \n",
+       "           3.425145   1.0      0.0           0.0      0.0    0.0         1.0   \n",
+       "           3.989158   1.0      0.0           0.0      1.0    0.0         1.0   \n",
+       "\n",
+       "                     serBilir   serChol   albumin  alkaline      SGOT  \\\n",
+       "sample_idx time_idx                                                     \n",
+       "1          0.569489  3.281890  0.000000 -0.894575  0.195532 -1.485263   \n",
+       "           1.095170  2.015877 -0.469461 -1.570646  0.285613  0.195488   \n",
+       "2          5.319790  0.172710 -0.658914 -1.431455 -0.605844 -0.442126   \n",
+       "           6.261636 -0.013468 -0.603657 -1.172958 -0.512364 -0.046806   \n",
+       "           7.266455  0.098239  0.000000 -1.312149 -0.443529  0.293680   \n",
+       "...                       ...       ...       ...       ...       ...   \n",
+       "312        1.045888  3.672865  3.319599  0.059878  1.385274  0.986129   \n",
+       "           1.867265  2.350998  2.901224 -0.099197  0.916176  0.641817   \n",
+       "           2.921367  0.694010 -0.066873  0.338261  0.327254  0.552551   \n",
+       "           3.425145  0.340271  0.000000 -0.377580  0.251620  0.016956   \n",
+       "           3.989158  0.507832  2.017110  0.795603  0.622990  0.169983   \n",
+       "\n",
+       "                     platelets  prothrombin       age  \n",
+       "sample_idx time_idx                                    \n",
+       "1          0.569489  -0.529101     0.136768  0.248058  \n",
+       "           1.095170  -0.456022     0.813132  0.248058  \n",
+       "2          5.319790  -1.395605     0.339677  1.292856  \n",
+       "           6.261636  -1.259888     0.339677  1.292856  \n",
+       "           7.266455  -1.364286     0.339677  1.292856  \n",
+       "...                        ...          ...       ...  \n",
+       "312        1.045888  -1.103291     1.624769 -1.962482  \n",
+       "           1.867265  -0.998892     1.354223 -1.962482  \n",
+       "           2.921367  -0.894494     0.474950 -1.962482  \n",
+       "           3.425145  -0.466462    -0.066141 -1.962482  \n",
+       "           3.989158  -0.351624    -0.133778 -1.962482  \n",
        "\n",
        "[1945 rows x 14 columns]"
       ]
@@ -2969,7 +3231,13 @@
       "\n",
       "PKPDDataSource loads the following dataset:\n",
       "\n",
-      "Generating simple PKPD dataset with random seed 100...\n",
+      "Generating simple PKPD dataset with random seed 100...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "OneOffTreatmentEffectsDataset(\n",
       "    time_series=TimeSeriesSamples([40, *, 2]),\n",
       "    predictive=OneOffTreatmentEffectsTaskData(\n",
@@ -3469,7 +3737,13 @@
       "--------------------------------------------------------------------------------\n",
       "\n",
       "DummyTemporalTreatmentEffectsDataSource loads the following dataset:\n",
-      "\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "TemporalTreatmentEffectsDataset(\n",
       "    time_series=TimeSeriesSamples([100, *, 5]),\n",
       "    static=StaticSamples([100, 3]),\n",
@@ -4096,8 +4370,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/data/tutorial04_data_splitting.ipynb b/tutorials/data/tutorial04_data_splitting.ipynb
index 504cbd05..36ed44af 100644
--- a/tutorials/data/tutorial04_data_splitting.ipynb
+++ b/tutorials/data/tutorial04_data_splitting.ipynb
@@ -39,7 +39,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -52,7 +52,7 @@
        ")"
       ]
      },
-     "execution_count": 1,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -77,7 +77,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -122,7 +122,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -145,7 +145,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -217,8 +217,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/extending/tutorial01_custom_method.ipynb b/tutorials/extending/tutorial01_custom_method.ipynb
index 5394964b..8b2909a3 100644
--- a/tutorials/extending/tutorial01_custom_method.ipynb
+++ b/tutorials/extending/tutorial01_custom_method.ipynb
@@ -61,7 +61,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -84,7 +84,7 @@
        " 'treatments.temporal.regression']"
       ]
      },
-     "execution_count": 18,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -107,7 +107,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -116,13 +116,7 @@
        "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'prediction'</span>: <span style=\"font-weight: bold\">{</span>\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span>\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span>\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'cde_classifier'</span>,\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'ode_classifier'</span>,\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'nn_classifier'</span>,\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'laplace_ode_classifier'</span>,\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'my_classifier'</span>\n",
-       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"font-weight: bold\">]</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'cde_classifier'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'ode_classifier'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'nn_classifier'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'laplace_ode_classifier'</span><span style=\"font-weight: bold\">]</span>,\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'laplace_ode_regressor'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'nn_regressor'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'ode_regressor'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'cde_regressor'</span><span style=\"font-weight: bold\">]</span>\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>,\n",
        "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'temporal'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'seq2seq_classifier'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'seq2seq_regressor'</span><span style=\"font-weight: bold\">]}</span>\n",
@@ -151,13 +145,7 @@
        "\u001b[1m{\u001b[0m\n",
        "\u001b[2;32m│   \u001b[0m\u001b[32m'prediction'\u001b[0m: \u001b[1m{\u001b[0m\n",
        "\u001b[2;32m│   │   \u001b[0m\u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\n",
-       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'cde_classifier'\u001b[0m,\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'ode_classifier'\u001b[0m,\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'nn_classifier'\u001b[0m,\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'laplace_ode_classifier'\u001b[0m,\n",
-       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'my_classifier'\u001b[0m\n",
-       "\u001b[2;32m│   │   │   \u001b[0m\u001b[1m]\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'cde_classifier'\u001b[0m, \u001b[32m'ode_classifier'\u001b[0m, \u001b[32m'nn_classifier'\u001b[0m, \u001b[32m'laplace_ode_classifier'\u001b[0m\u001b[1m]\u001b[0m,\n",
        "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'laplace_ode_regressor'\u001b[0m, \u001b[32m'nn_regressor'\u001b[0m, \u001b[32m'ode_regressor'\u001b[0m, \u001b[32m'cde_regressor'\u001b[0m\u001b[1m]\u001b[0m\n",
        "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m,\n",
        "\u001b[2;32m│   │   \u001b[0m\u001b[32m'temporal'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'seq2seq_classifier'\u001b[0m\u001b[1m]\u001b[0m, \u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'seq2seq_regressor'\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\n",
@@ -206,7 +194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -417,7 +405,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -480,7 +468,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -588,7 +576,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -614,7 +602,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -721,7 +709,7 @@
        "[100 rows x 1 columns]"
       ]
      },
-     "execution_count": 24,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -778,8 +766,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/extending/tutorial03_custom_datasource.ipynb b/tutorials/extending/tutorial03_custom_datasource.ipynb
index 4bedfb51..4ce6f99a 100644
--- a/tutorials/extending/tutorial03_custom_datasource.ipynb
+++ b/tutorials/extending/tutorial03_custom_datasource.ipynb
@@ -53,7 +53,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -66,7 +66,7 @@
        " 'treatments.temporal']"
       ]
      },
-     "execution_count": 1,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -84,12 +84,12 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Remember you can also see the existing method plugins and how they correspond to different categories, as follows:"
+    "Remember you can also see the existing data source plugins and how they correspond to different categories, as follows:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -134,7 +134,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -249,7 +249,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -292,7 +292,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -352,14 +352,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "<__main__.MyDataSource object at 0x7fa2cf032850>\n"
+      "<__main__.MyDataSource object at 0x7f5f36a6cbb0>\n"
      ]
     }
    ],
@@ -373,7 +373,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -385,7 +385,7 @@
        ")"
       ]
      },
-     "execution_count": 7,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -400,7 +400,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -636,7 +636,7 @@
        "[3000 rows x 10 columns]"
       ]
      },
-     "execution_count": 8,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -649,7 +649,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -749,7 +749,7 @@
        "[100 rows x 1 columns]"
       ]
      },
-     "execution_count": 9,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -801,8 +801,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/extending/tutorial04_custom_metric.ipynb b/tutorials/extending/tutorial04_custom_metric.ipynb
new file mode 100644
index 00000000..b7a05185
--- /dev/null
+++ b/tutorials/extending/tutorial04_custom_metric.ipynb
@@ -0,0 +1,505 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Extending TemporAI Tutorial 04: Writing a Custom Metric Plugin\n",
+    "\n",
+    "This tutorial shows how to extend TemporAI by wring a custom *metric* plugin."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "*Skip the below cell if you are not on Google Colab / already have TemporAI installed:*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install temporai\n",
+    "\n",
+    "# Or from the repo, for the latest version:\n",
+    "# %pip install git+https://github.com/vanderschaarlab/temporai.git"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Note**\n",
+    "\n",
+    "See also \"Writing a Custom `Plugin` 101\" section in \"Writing a Custom Method Plugin\" tutorial."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Inherit from the appropriate **base class** for the category of the metric plugin you are writing.\n",
+    "\n",
+    "You need to find which category of metric plugin you are writing.\n",
+    "\n",
+    "You can view all the different metric plugin categories as so:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['prediction.one_off.classification',\n",
+       " 'prediction.one_off.regression',\n",
+       " 'time_to_event']"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from tempor import plugin_loader\n",
+    "\n",
+    "plugin_categories = plugin_loader.list_categories(plugin_type=\"metric\")\n",
+    "\n",
+    "list(plugin_categories.keys())"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Remember you can also see the existing metric plugins and how they correspond to different categories, as follows:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'prediction'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'mse'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'mae'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'r2'</span><span style=\"font-weight: bold\">]</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'accuracy'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'kappa'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'kappa_quadratic'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'mcc'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'aucprc'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'aucroc'</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"font-weight: bold\">]</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"font-weight: bold\">}</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'time_to_event'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'c_index'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'brier_score'</span><span style=\"font-weight: bold\">]</span>\n",
+       "<span style=\"font-weight: bold\">}</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[32m'prediction'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   │   \u001b[0m\u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'mse'\u001b[0m, \u001b[32m'mae'\u001b[0m, \u001b[32m'r2'\u001b[0m\u001b[1m]\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'accuracy'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'kappa'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'kappa_quadratic'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'mcc'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'aucprc'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'aucroc'\u001b[0m\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[1m]\u001b[0m\n",
+       "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[1m}\u001b[0m,\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[32m'time_to_event'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'c_index'\u001b[0m, \u001b[32m'brier_score'\u001b[0m\u001b[1m]\u001b[0m\n",
+       "\u001b[1m}\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "all_plugins = plugin_loader.list(plugin_type=\"metric\")\n",
+    "\n",
+    "from rich.pretty import pprint  # For prettifying the print output only.\n",
+    "\n",
+    "pprint(all_plugins, indent_guides=True)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's say you would like to write a plugin of category `\"prediction.one_off.regression\"`.\n",
+    "\n",
+    "You can find which base class you need to inherit from as follows."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Base classes for all categories:\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
+       "    <span style=\"color: #008000; text-decoration-color: #008000\">'prediction.one_off.classification'</span>: <span style=\"font-weight: bold\">&lt;</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff; font-weight: bold\">class</span><span style=\"color: #000000; text-decoration-color: #000000\"> </span><span style=\"color: #008000; text-decoration-color: #008000\">'tempor.metrics.metric.OneOffClassificationMetric'</span><span style=\"color: #000000; text-decoration-color: #000000\">&gt;,</span>\n",
+       "<span style=\"color: #000000; text-decoration-color: #000000\">    </span><span style=\"color: #008000; text-decoration-color: #008000\">'prediction.one_off.regression'</span><span style=\"color: #000000; text-decoration-color: #000000\">: &lt;class </span><span style=\"color: #008000; text-decoration-color: #008000\">'tempor.metrics.metric.OneOffRegressionMetric'</span><span style=\"color: #000000; text-decoration-color: #000000\">&gt;,</span>\n",
+       "<span style=\"color: #000000; text-decoration-color: #000000\">    </span><span style=\"color: #008000; text-decoration-color: #008000\">'time_to_event'</span><span style=\"color: #000000; text-decoration-color: #000000\">: &lt;class </span><span style=\"color: #008000; text-decoration-color: #008000\">'tempor.metrics.metric.TimeToEventMetric'</span><span style=\"font-weight: bold\">&gt;</span>\n",
+       "<span style=\"font-weight: bold\">}</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[1m{\u001b[0m\n",
+       "    \u001b[32m'prediction.one_off.classification'\u001b[0m: \u001b[1m<\u001b[0m\u001b[1;95mclass\u001b[0m\u001b[39m \u001b[0m\u001b[32m'tempor.metrics.metric.OneOffClassificationMetric'\u001b[0m\u001b[39m>,\u001b[0m\n",
+       "\u001b[39m    \u001b[0m\u001b[32m'prediction.one_off.regression'\u001b[0m\u001b[39m: <class \u001b[0m\u001b[32m'tempor.metrics.metric.OneOffRegressionMetric'\u001b[0m\u001b[39m>,\u001b[0m\n",
+       "\u001b[39m    \u001b[0m\u001b[32m'time_to_event'\u001b[0m\u001b[39m: <class \u001b[0m\u001b[32m'tempor.metrics.metric.TimeToEventMetric'\u001b[0m\u001b[1m>\u001b[0m\n",
+       "\u001b[1m}\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Base class you need:\n",
+      "<class 'tempor.metrics.metric.OneOffRegressionMetric'>\n"
+     ]
+    }
+   ],
+   "source": [
+    "plugin_categories = plugin_loader.list_categories(plugin_type=\"metric\")\n",
+    "\n",
+    "print(\"Base classes for all categories:\")\n",
+    "pprint(plugin_categories, indent_guides=False)\n",
+    "\n",
+    "print(\"Base class you need:\")\n",
+    "print(plugin_categories[\"prediction.one_off.regression\"])"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can then find the class in the TemporAI source code, to see its method signatures etc."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Implement the **methods** the plugin needs.\n",
+    "\n",
+    "`Metric` plugins require the following methods to be implemented:\n",
+    "* `direction` property, which returns either `\"minimize\"` or `\"maximize\"`, representing the \"good\" direction of the metric.\n",
+    "* `_evaluate()` which takes in the actual and predicted values and returns the evaluated metric(s).\n",
+    "\n",
+    "If you haven't implemented some required method for the plugin, Python will notify you by raising an exception when you\n",
+    "attempt to instantiate your plugin (see [Python `abc`](https://docs.python.org/3/library/abc.html)).\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Register** the plugin with TemporAI.\n",
+    "\n",
+    "Registering your plugin with TemporAI is very simple, you need to use the `register_plugin` decorator,\n",
+    "as shown in the example below.\n",
+    "\n",
+    "You will need to specify the `name` of your plugin and its `category` in the decorator.\n",
+    "\n",
+    "The `plugin_type` needs to be set to `\"datasource\"`.\n",
+    "\n",
+    "```python\n",
+    "from tempor.core.plugins import register_plugin\n",
+    "\n",
+    "@register_plugin(name=\"my_metric\", category=\"prediction.one_off.regression\", plugin_type=\"metric\")\n",
+    "class MyMetric(OneOffRegressionMetric):\n",
+    "    ...\n",
+    "```"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Example\n",
+    "\n",
+    "Now putting this together in a minimal example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Any\n",
+    "\n",
+    "import numpy as np\n",
+    "from sklearn.metrics import d2_tweedie_score\n",
+    "\n",
+    "from tempor.core import plugins\n",
+    "from tempor.metrics import metric, metric_typing\n",
+    "\n",
+    "\n",
+    "@plugins.register_plugin(name=\"my_metric\", category=\"prediction.one_off.regression\", plugin_type=\"metric\")\n",
+    "class MyOneOffRegressionMetric(metric.OneOffRegressionMetric):\n",
+    "    \"\"\"My custom metric, here we use\n",
+    "    D^2 regression score function, fraction of Tweedie deviance explained,\n",
+    "    as per `sklearn`.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    @property\n",
+    "    def direction(self) -> metric_typing.MetricDirection:  # noqa: D102\n",
+    "        return \"maximize\"\n",
+    "\n",
+    "    def _evaluate(self, actual: np.ndarray, predicted: np.ndarray, *args: Any, **kwargs: Any) -> float:\n",
+    "        return d2_tweedie_score(actual, predicted, power=kwargs.get(\"power\", 0))"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We now see our plugin in TemporAI:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'prediction'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'mse'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'mae'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'r2'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'my_metric'</span><span style=\"font-weight: bold\">]</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'accuracy'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'kappa'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'kappa_quadratic'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'recall_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_micro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_macro'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'precision_weighted'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'mcc'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'aucprc'</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'aucroc'</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"font-weight: bold\">]</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"font-weight: bold\">}</span>,\n",
+       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'time_to_event'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'c_index'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'brier_score'</span><span style=\"font-weight: bold\">]</span>\n",
+       "<span style=\"font-weight: bold\">}</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[32m'prediction'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   │   \u001b[0m\u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'mse'\u001b[0m, \u001b[32m'mae'\u001b[0m, \u001b[32m'r2'\u001b[0m, \u001b[32m'my_metric'\u001b[0m\u001b[1m]\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'accuracy'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'f1_score_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'kappa'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'kappa_quadratic'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'recall_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_micro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_macro'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'precision_weighted'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'mcc'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'aucprc'\u001b[0m,\n",
+       "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[32m'aucroc'\u001b[0m\n",
+       "\u001b[2;32m│   │   │   \u001b[0m\u001b[1m]\u001b[0m\n",
+       "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[1m}\u001b[0m,\n",
+       "\u001b[2;32m│   \u001b[0m\u001b[32m'time_to_event'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'c_index'\u001b[0m, \u001b[32m'brier_score'\u001b[0m\u001b[1m]\u001b[0m\n",
+       "\u001b[1m}\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "`my_metric` plugin found in the category 'prediction.one_off.regression': True\n"
+     ]
+    }
+   ],
+   "source": [
+    "from tempor import plugin_loader\n",
+    "\n",
+    "all_plugins = plugin_loader.list(plugin_type=\"metric\")\n",
+    "\n",
+    "pprint(all_plugins, indent_guides=True)\n",
+    "\n",
+    "my_metric_found = \"my_metric\" in all_plugins[\"prediction\"][\"one_off\"][\"regression\"]\n",
+    "print(f\"`my_metric` plugin found in the category 'prediction.one_off.regression': {my_metric_found}\")\n",
+    "assert my_metric_found"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The plugin can be used as normal."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "MyOneOffRegressionMetric(\n",
+      "    name='my_metric',\n",
+      "    description='My custom metric, here we use\\n    D^2 regression score function, fraction of Tweedie deviance explai...'\n",
+      ")\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Get the plugin.\n",
+    "\n",
+    "my_metric = plugin_loader.get(\"prediction.one_off.regression.my_metric\", plugin_type=\"metric\")\n",
+    "\n",
+    "print(my_metric)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.7492656592741931\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Use the metric.\n",
+    "\n",
+    "y_true = np.asarray([1, 2, 3])\n",
+    "y_pred = np.asarray([1, 3, 3])\n",
+    "\n",
+    "metric_value = my_metric(y_true, y_pred, power=2)\n",
+    "print(metric_value)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 🎉 Congratulations!\n",
+    "\n",
+    "Congratulations on completing this notebook tutorial! If you enjoyed this and would like to join the movement towards *Machine learning and AI for Medicine*, you can do so in the following ways!\n",
+    "\n",
+    "\n",
+    "\n",
+    "### ⭐ Star [TemporAI](https://github.com/vanderschaarlab/temporai) on GitHub\n",
+    "\n",
+    "- The easiest way to help our community is by just starring the repos! This helps raise awareness of the tools we're building.\n",
+    "\n",
+    "\n",
+    "\n",
+    "### Check out other projects from [vanderschaarlab](https://github.com/vanderschaarlab)\n",
+    "- 📝 [HyperImpute](https://github.com/vanderschaarlab/hyperimpute)\n",
+    "- 📊 [AutoPrognosis](https://github.com/vanderschaarlab/autoprognosis)\n",
+    "- 🤖 [SynthCity](https://github.com/vanderschaarlab/synthcity)\n",
+    " "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "RE_py38_temporai",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/tutorials/usage/tutorial01_plugins.ipynb b/tutorials/usage/tutorial01_plugins.ipynb
index 6e2841c4..416e04b5 100644
--- a/tutorials/usage/tutorial01_plugins.ipynb
+++ b/tutorials/usage/tutorial01_plugins.ipynb
@@ -50,7 +50,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -113,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -188,7 +188,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -281,7 +281,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -316,7 +316,7 @@
        " 'treatments.temporal.regression.crn_regressor']"
       ]
      },
-     "execution_count": 10,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -346,16 +346,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "['datasource', 'method']"
+       "['datasource', 'method', 'metric']"
       ]
      },
-     "execution_count": 11,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -373,7 +373,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -410,6 +410,31 @@
        "            <span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'synctwin_regressor'</span><span style=\"font-weight: bold\">]}</span>,\n",
        "            <span style=\"color: #008000; text-decoration-color: #008000\">'temporal'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'crn_classifier'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'crn_regressor'</span><span style=\"font-weight: bold\">]}</span>\n",
        "        <span style=\"font-weight: bold\">}</span>\n",
+       "    <span style=\"font-weight: bold\">}</span>,\n",
+       "    <span style=\"color: #008000; text-decoration-color: #008000\">'metric'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'prediction'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "            <span style=\"color: #008000; text-decoration-color: #008000\">'one_off'</span>: <span style=\"font-weight: bold\">{</span>\n",
+       "                <span style=\"color: #008000; text-decoration-color: #008000\">'regression'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'mse'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'mae'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'r2'</span><span style=\"font-weight: bold\">]</span>,\n",
+       "                <span style=\"color: #008000; text-decoration-color: #008000\">'classification'</span>: <span style=\"font-weight: bold\">[</span>\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'accuracy'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_micro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_macro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'f1_score_weighted'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'kappa'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'kappa_quadratic'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'recall_micro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'recall_macro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'recall_weighted'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'precision_micro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'precision_macro'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'precision_weighted'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'mcc'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'aucprc'</span>,\n",
+       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'aucroc'</span>\n",
+       "                <span style=\"font-weight: bold\">]</span>\n",
+       "            <span style=\"font-weight: bold\">}</span>\n",
+       "        <span style=\"font-weight: bold\">}</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'time_to_event'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'c_index'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'brier_score'</span><span style=\"font-weight: bold\">]</span>\n",
        "    <span style=\"font-weight: bold\">}</span>\n",
        "<span style=\"font-weight: bold\">}</span>\n",
        "</pre>\n"
@@ -446,6 +471,31 @@
        "            \u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'synctwin_regressor'\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m,\n",
        "            \u001b[32m'temporal'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'crn_classifier'\u001b[0m\u001b[1m]\u001b[0m, \u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'crn_regressor'\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\n",
        "        \u001b[1m}\u001b[0m\n",
+       "    \u001b[1m}\u001b[0m,\n",
+       "    \u001b[32m'metric'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "        \u001b[32m'prediction'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "            \u001b[32m'one_off'\u001b[0m: \u001b[1m{\u001b[0m\n",
+       "                \u001b[32m'regression'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'mse'\u001b[0m, \u001b[32m'mae'\u001b[0m, \u001b[32m'r2'\u001b[0m\u001b[1m]\u001b[0m,\n",
+       "                \u001b[32m'classification'\u001b[0m: \u001b[1m[\u001b[0m\n",
+       "                    \u001b[32m'accuracy'\u001b[0m,\n",
+       "                    \u001b[32m'f1_score_micro'\u001b[0m,\n",
+       "                    \u001b[32m'f1_score_macro'\u001b[0m,\n",
+       "                    \u001b[32m'f1_score_weighted'\u001b[0m,\n",
+       "                    \u001b[32m'kappa'\u001b[0m,\n",
+       "                    \u001b[32m'kappa_quadratic'\u001b[0m,\n",
+       "                    \u001b[32m'recall_micro'\u001b[0m,\n",
+       "                    \u001b[32m'recall_macro'\u001b[0m,\n",
+       "                    \u001b[32m'recall_weighted'\u001b[0m,\n",
+       "                    \u001b[32m'precision_micro'\u001b[0m,\n",
+       "                    \u001b[32m'precision_macro'\u001b[0m,\n",
+       "                    \u001b[32m'precision_weighted'\u001b[0m,\n",
+       "                    \u001b[32m'mcc'\u001b[0m,\n",
+       "                    \u001b[32m'aucprc'\u001b[0m,\n",
+       "                    \u001b[32m'aucroc'\u001b[0m\n",
+       "                \u001b[1m]\u001b[0m\n",
+       "            \u001b[1m}\u001b[0m\n",
+       "        \u001b[1m}\u001b[0m,\n",
+       "        \u001b[32m'time_to_event'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'c_index'\u001b[0m, \u001b[32m'brier_score'\u001b[0m\u001b[1m]\u001b[0m\n",
        "    \u001b[1m}\u001b[0m\n",
        "\u001b[1m}\u001b[0m\n"
       ]
@@ -503,8 +553,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/usage/tutorial02_imputation.ipynb b/tutorials/usage/tutorial02_imputation.ipynb
index 95daaaeb..3ea3c5bc 100644
--- a/tutorials/usage/tutorial02_imputation.ipynb
+++ b/tutorials/usage/tutorial02_imputation.ipynb
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -51,7 +51,7 @@
        " 'temporal': ['ffill', 'ts_tabular_imputer', 'bfill']}"
       ]
      },
-     "execution_count": 1,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -71,7 +71,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,14 +88,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
      ]
     },
     {
@@ -132,7 +132,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -278,7 +278,7 @@
        "[100 rows x 4 columns]"
       ]
      },
-     "execution_count": 4,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -293,48 +293,252 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.22005514968324613 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21670750510884584 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.2166465658117811 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.2166456843686172 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:46:52 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 7\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 8\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 9\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 10\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n",
-      "2023-10-10 16:46:53 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.22005514968324613 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21670750510884584 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.2166465658117811 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.2166456843686172 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:42 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 7\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 8\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 9\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 10\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 0 <-- score -0.23109777030617995 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>> Column 1 <-- score -0.21664567863448503 <-- Model linear_regression\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
      ]
     },
     {
@@ -480,7 +684,7 @@
        "[100 rows x 4 columns]"
       ]
      },
-     "execution_count": 5,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -505,7 +709,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -538,7 +742,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -702,7 +906,7 @@
        "[1000 rows x 5 columns]"
       ]
      },
-     "execution_count": 7,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -717,7 +921,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -881,7 +1085,7 @@
        "[1000 rows x 5 columns]"
       ]
      },
-     "execution_count": 8,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -937,8 +1141,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/usage/tutorial03_scaling.ipynb b/tutorials/usage/tutorial03_scaling.ipynb
index 1c5c2534..519fd880 100644
--- a/tutorials/usage/tutorial03_scaling.ipynb
+++ b/tutorials/usage/tutorial03_scaling.ipynb
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -51,7 +51,7 @@
        " 'temporal': ['ts_minmax_scaler', 'ts_standard_scaler']}"
       ]
      },
-     "execution_count": 1,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -71,7 +71,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,7 +88,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -121,7 +121,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -328,7 +328,7 @@
        "[100 rows x 4 columns]"
       ]
      },
-     "execution_count": 4,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -346,7 +346,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -553,7 +553,7 @@
        "[100 rows x 4 columns]"
       ]
      },
-     "execution_count": 5,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -579,7 +579,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -612,7 +612,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -843,7 +843,7 @@
        "[1000 rows x 5 columns]"
       ]
      },
-     "execution_count": 7,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -861,7 +861,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1092,7 +1092,7 @@
        "[1000 rows x 5 columns]"
       ]
      },
-     "execution_count": 8,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1149,8 +1149,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/usage/tutorial04_prediction.ipynb b/tutorials/usage/tutorial04_prediction.ipynb
index 46ad468e..ae2299ab 100644
--- a/tutorials/usage/tutorial04_prediction.ipynb
+++ b/tutorials/usage/tutorial04_prediction.ipynb
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -88,7 +88,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -110,7 +110,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -163,7 +163,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -263,7 +263,7 @@
        "[100 rows x 1 columns]"
       ]
      },
-     "execution_count": 4,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -275,18 +275,42 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.68824702501297, validation loss: 0.6883962154388428\n",
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:10| train loss: 0.6999298334121704, validation loss: 0.6879696846008301\n",
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:20| train loss: 0.6818332672119141, validation loss: 0.6867876648902893\n",
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:30| train loss: 0.6943979263305664, validation loss: 0.685184121131897\n",
-      "2023-10-10 16:49:25 | INFO     | tempor.models.ts_model:_train:379 | Epoch:40| train loss: 0.6284184455871582, validation loss: 0.6934055089950562\n"
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.68824702501297, validation loss: 0.6883962154388428\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:10| train loss: 0.6999298930168152, validation loss: 0.6879696249961853\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:20| train loss: 0.6818333268165588, validation loss: 0.6867876052856445\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:30| train loss: 0.694374680519104, validation loss: 0.6851812601089478\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:40| train loss: 0.6282318234443665, validation loss: 0.6936542391777039\n"
      ]
     },
     {
@@ -320,7 +344,7 @@
        ")"
       ]
      },
-     "execution_count": 5,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -332,7 +356,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -432,7 +456,7 @@
        "[100 rows x 1 columns]"
       ]
      },
-     "execution_count": 6,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -455,7 +479,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -517,7 +541,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -635,7 +659,7 @@
        "[1573 rows x 2 columns]"
       ]
      },
-     "execution_count": 8,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -647,37 +671,121 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Preparing data for decoder training...\n",
-      "Preparing data for decoder training DONE.\n",
+      "Preparing data for decoder training...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Preparing data for decoder training DONE.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "=== Training stage: 1. Train encoder ===\n",
       "Epoch: 0, Loss: 68.64124755859375\n",
-      "Epoch: 1, Loss: 35.37634651184082\n",
-      "Epoch: 2, Loss: 21.269977149963378\n",
-      "Epoch: 3, Loss: 17.418054847717286\n",
-      "Epoch: 4, Loss: 13.851956520080567\n",
-      "Epoch: 5, Loss: 8.335217247009277\n",
-      "Epoch: 6, Loss: 6.385376853942871\n",
-      "Epoch: 7, Loss: 5.112505340576172\n",
-      "Epoch: 8, Loss: 4.690582332611084\n",
-      "Epoch: 9, Loss: 4.158801422119141\n",
-      "=== Training stage: 2. Train decoder ===\n",
-      "Epoch: 0, Loss: 29.110631188485222\n",
-      "Epoch: 1, Loss: 4.417614046603097\n",
-      "Epoch: 2, Loss: 3.908697671333463\n",
-      "Epoch: 3, Loss: 3.8485680392058788\n",
-      "Epoch: 4, Loss: 3.8276885692068126\n",
-      "Epoch: 5, Loss: 3.793576312405029\n",
-      "Epoch: 6, Loss: 3.749186389072364\n",
-      "Epoch: 7, Loss: 3.787092407949102\n",
-      "Epoch: 8, Loss: 3.7378030509586164\n",
-      "Epoch: 9, Loss: 3.717736062595655\n"
+      "Epoch: 1, Loss: 35.376345291137696\n",
+      "Epoch: 2, Loss: 21.269977149963378\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 3, Loss: 17.418056030273437\n",
+      "Epoch: 4, Loss: 13.851958045959472\n",
+      "Epoch: 5, Loss: 8.335217628479004\n",
+      "Epoch: 6, Loss: 6.385376873016358\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 7, Loss: 5.1125048828125\n",
+      "Epoch: 8, Loss: 4.690582141876221\n",
+      "Epoch: 9, Loss: 4.15880220413208\n",
+      "=== Training stage: 2. Train decoder ===\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 0, Loss: 29.110622051453316\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 1, Loss: 4.4235033275153635\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 2, Loss: 3.954428958957566\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 3, Loss: 3.8989672662446226\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 4, Loss: 3.8317033024119103\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 5, Loss: 3.8341090269664964\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 6, Loss: 3.717077842797864\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 7, Loss: 3.7477998285170564\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 8, Loss: 3.859872080239\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 9, Loss: 3.8217773019903305\n"
      ]
     }
    ],
@@ -688,7 +796,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -727,28 +835,28 @@
        "    <tr>\n",
        "      <th rowspan=\"5\" valign=\"top\">0</th>\n",
        "      <th>11</th>\n",
-       "      <td>9.828533</td>\n",
-       "      <td>8.405809</td>\n",
+       "      <td>10.618200</td>\n",
+       "      <td>8.973257</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>12</th>\n",
-       "      <td>9.831676</td>\n",
-       "      <td>8.414574</td>\n",
+       "      <td>10.728621</td>\n",
+       "      <td>9.108867</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>13</th>\n",
-       "      <td>9.813161</td>\n",
-       "      <td>8.395016</td>\n",
+       "      <td>10.735380</td>\n",
+       "      <td>9.118745</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>14</th>\n",
-       "      <td>9.798813</td>\n",
-       "      <td>8.378744</td>\n",
+       "      <td>10.728010</td>\n",
+       "      <td>9.093669</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>15</th>\n",
-       "      <td>9.772892</td>\n",
-       "      <td>8.352385</td>\n",
+       "      <td>10.738855</td>\n",
+       "      <td>9.097306</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>...</th>\n",
@@ -759,28 +867,28 @@
        "    <tr>\n",
        "      <th rowspan=\"5\" valign=\"top\">99</th>\n",
        "      <th>17</th>\n",
-       "      <td>11.214181</td>\n",
-       "      <td>9.754721</td>\n",
+       "      <td>12.021331</td>\n",
+       "      <td>10.417010</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>18</th>\n",
-       "      <td>11.275115</td>\n",
-       "      <td>9.821613</td>\n",
+       "      <td>12.091549</td>\n",
+       "      <td>10.484824</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>19</th>\n",
-       "      <td>11.298031</td>\n",
-       "      <td>9.848743</td>\n",
+       "      <td>12.144796</td>\n",
+       "      <td>10.523143</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>20</th>\n",
-       "      <td>11.280746</td>\n",
-       "      <td>9.832634</td>\n",
+       "      <td>12.133350</td>\n",
+       "      <td>10.510198</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>21</th>\n",
-       "      <td>11.256371</td>\n",
-       "      <td>9.801851</td>\n",
+       "      <td>12.097664</td>\n",
+       "      <td>10.484853</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -789,24 +897,24 @@
       ],
       "text/plain": [
        "TimeSeriesSamples with data:\n",
-       "                             0         1\n",
-       "sample_idx time_idx                     \n",
-       "0          11         9.828533  8.405809\n",
-       "           12         9.831676  8.414574\n",
-       "           13         9.813161  8.395016\n",
-       "           14         9.798813  8.378744\n",
-       "           15         9.772892  8.352385\n",
-       "...                        ...       ...\n",
-       "99         17        11.214181  9.754721\n",
-       "           18        11.275115  9.821613\n",
-       "           19        11.298031  9.848743\n",
-       "           20        11.280746  9.832634\n",
-       "           21        11.256371  9.801851\n",
+       "                             0          1\n",
+       "sample_idx time_idx                      \n",
+       "0          11        10.618200   8.973257\n",
+       "           12        10.728621   9.108867\n",
+       "           13        10.735380   9.118745\n",
+       "           14        10.728010   9.093669\n",
+       "           15        10.738855   9.097306\n",
+       "...                        ...        ...\n",
+       "99         17        12.021331  10.417010\n",
+       "           18        12.091549  10.484824\n",
+       "           19        12.144796  10.523143\n",
+       "           20        12.133350  10.510198\n",
+       "           21        12.097664  10.484853\n",
        "\n",
        "[1000 rows x 2 columns]"
       ]
      },
-     "execution_count": 10,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -858,8 +966,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/usage/tutorial05_time_to_event.ipynb b/tutorials/usage/tutorial05_time_to_event.ipynb
index aab5b865..9ed752d1 100644
--- a/tutorials/usage/tutorial05_time_to_event.ipynb
+++ b/tutorials/usage/tutorial05_time_to_event.ipynb
@@ -44,7 +44,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -53,7 +53,7 @@
        "['ts_coxph', 'ts_xgb', 'dynamic_deephit']"
       ]
      },
-     "execution_count": 1,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -73,7 +73,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -90,7 +90,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -140,7 +140,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -240,7 +240,7 @@
        "[312 rows x 1 columns]"
       ]
      },
-     "execution_count": 4,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -252,7 +252,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -283,7 +283,7 @@
        ")"
       ]
      },
-     "execution_count": 5,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -295,7 +295,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -332,24 +332,24 @@
        "    <tr>\n",
        "      <th rowspan=\"3\" valign=\"top\">1</th>\n",
        "      <th>0.25</th>\n",
-       "      <td>0.422643</td>\n",
+       "      <td>0.422638</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.50</th>\n",
-       "      <td>0.626620</td>\n",
+       "      <td>0.626617</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.75</th>\n",
-       "      <td>0.717206</td>\n",
+       "      <td>0.717204</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th rowspan=\"2\" valign=\"top\">2</th>\n",
        "      <th>0.25</th>\n",
-       "      <td>0.221314</td>\n",
+       "      <td>0.221283</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.50</th>\n",
-       "      <td>0.288492</td>\n",
+       "      <td>0.288457</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>...</th>\n",
@@ -359,24 +359,24 @@
        "    <tr>\n",
        "      <th rowspan=\"2\" valign=\"top\">311</th>\n",
        "      <th>0.50</th>\n",
-       "      <td>0.039302</td>\n",
+       "      <td>0.039307</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.75</th>\n",
-       "      <td>0.048346</td>\n",
+       "      <td>0.048353</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th rowspan=\"3\" valign=\"top\">312</th>\n",
        "      <th>0.25</th>\n",
-       "      <td>0.248834</td>\n",
+       "      <td>0.248848</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.50</th>\n",
-       "      <td>0.390042</td>\n",
+       "      <td>0.390068</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>0.75</th>\n",
-       "      <td>0.444533</td>\n",
+       "      <td>0.444567</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -387,22 +387,22 @@
        "TimeSeriesSamples with data:\n",
        "                     risk_score\n",
        "sample_idx time_idx            \n",
-       "1          0.25        0.422643\n",
-       "           0.50        0.626620\n",
-       "           0.75        0.717206\n",
-       "2          0.25        0.221314\n",
-       "           0.50        0.288492\n",
+       "1          0.25        0.422638\n",
+       "           0.50        0.626617\n",
+       "           0.75        0.717204\n",
+       "2          0.25        0.221283\n",
+       "           0.50        0.288457\n",
        "...                         ...\n",
-       "311        0.50        0.039302\n",
-       "           0.75        0.048346\n",
-       "312        0.25        0.248834\n",
-       "           0.50        0.390042\n",
-       "           0.75        0.444533\n",
+       "311        0.50        0.039307\n",
+       "           0.75        0.048353\n",
+       "312        0.25        0.248848\n",
+       "           0.50        0.390068\n",
+       "           0.75        0.444567\n",
        "\n",
        "[936 rows x 1 columns]"
       ]
      },
-     "execution_count": 6,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -466,8 +466,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/usage/tutorial06_treatments.ipynb b/tutorials/usage/tutorial06_treatments.ipynb
index bb7590cf..11e4efa0 100644
--- a/tutorials/usage/tutorial06_treatments.ipynb
+++ b/tutorials/usage/tutorial06_treatments.ipynb
@@ -45,7 +45,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -86,7 +86,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -107,7 +107,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -175,7 +175,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -293,7 +293,7 @@
        "[1573 rows x 2 columns]"
       ]
      },
-     "execution_count": 5,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -305,7 +305,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -410,7 +410,7 @@
        "[1573 rows x 1 columns]"
       ]
      },
-     "execution_count": 6,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -422,57 +422,225 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Preparing data for decoder training...\n",
-      "Preparing data for decoder training DONE.\n",
+      "Preparing data for decoder training...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Preparing data for decoder training DONE.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "=== Training stage: 1. Train encoder ===\n",
-      "Epoch: 0, Prediction Loss: 75.327, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 75.327\n",
-      "Epoch: 1, Prediction Loss: 37.270, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 37.270\n",
-      "Epoch: 2, Prediction Loss: 21.899, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 21.899\n",
-      "Epoch: 3, Prediction Loss: 18.649, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 18.649\n",
-      "Epoch: 4, Prediction Loss: 18.343, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 18.343\n",
-      "Epoch: 5, Prediction Loss: 15.802, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 15.802\n",
-      "Epoch: 6, Prediction Loss: 10.720, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 10.720\n",
-      "Epoch: 7, Prediction Loss: 8.111, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 8.111\n",
-      "Epoch: 8, Prediction Loss: 6.840, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 6.840\n",
-      "Epoch: 9, Prediction Loss: 5.708, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.708\n",
-      "Epoch: 10, Prediction Loss: 5.061, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.061\n",
-      "Epoch: 11, Prediction Loss: 4.648, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.648\n",
-      "Epoch: 12, Prediction Loss: 4.420, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.420\n",
-      "Epoch: 13, Prediction Loss: 4.356, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.356\n",
-      "Epoch: 14, Prediction Loss: 4.175, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.175\n",
-      "Epoch: 15, Prediction Loss: 4.014, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.014\n",
-      "Epoch: 16, Prediction Loss: 3.899, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.899\n",
-      "Epoch: 17, Prediction Loss: 3.933, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.933\n",
-      "Epoch: 18, Prediction Loss: 3.940, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.940\n",
-      "Epoch: 19, Prediction Loss: 3.844, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.844\n",
-      "=== Training stage: 2. Train decoder ===\n",
-      "Epoch: 0, Prediction Loss: 34.517, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 34.517\n",
-      "Epoch: 1, Prediction Loss: 7.175, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 7.175\n",
-      "Epoch: 2, Prediction Loss: 3.985, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.985\n",
-      "Epoch: 3, Prediction Loss: 3.831, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.831\n",
-      "Epoch: 4, Prediction Loss: 3.851, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.851\n",
-      "Epoch: 5, Prediction Loss: 3.730, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.730\n",
-      "Epoch: 6, Prediction Loss: 3.767, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.767\n",
-      "Epoch: 7, Prediction Loss: 3.737, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.737\n",
-      "Epoch: 8, Prediction Loss: 3.678, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.678\n",
-      "Epoch: 9, Prediction Loss: 3.659, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.659\n",
-      "Epoch: 10, Prediction Loss: 3.698, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.698\n",
-      "Epoch: 11, Prediction Loss: 3.684, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.684\n",
-      "Epoch: 12, Prediction Loss: 3.678, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.678\n",
-      "Epoch: 13, Prediction Loss: 3.649, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.649\n",
-      "Epoch: 14, Prediction Loss: 3.617, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.617\n",
-      "Epoch: 15, Prediction Loss: 3.670, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.670\n",
-      "Epoch: 16, Prediction Loss: 3.638, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.638\n",
-      "Epoch: 17, Prediction Loss: 3.655, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.655\n",
-      "Epoch: 18, Prediction Loss: 3.654, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.654\n",
-      "Epoch: 19, Prediction Loss: 3.668, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.668\n"
+      "Epoch: 0, Prediction Loss: 75.212, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 75.212\n",
+      "Epoch: 1, Prediction Loss: 32.453, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 32.453\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 2, Prediction Loss: 19.389, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 19.389\n",
+      "Epoch: 3, Prediction Loss: 19.679, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 19.679\n",
+      "Epoch: 4, Prediction Loss: 19.588, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 19.588\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 5, Prediction Loss: 14.915, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 14.915\n",
+      "Epoch: 6, Prediction Loss: 10.608, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 10.608\n",
+      "Epoch: 7, Prediction Loss: 8.684, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 8.684\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 8, Prediction Loss: 6.953, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 6.953\n",
+      "Epoch: 9, Prediction Loss: 5.645, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.645\n",
+      "Epoch: 10, Prediction Loss: 5.060, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.060\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 11, Prediction Loss: 4.658, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.658\n",
+      "Epoch: 12, Prediction Loss: 4.398, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.398\n",
+      "Epoch: 13, Prediction Loss: 4.265, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.265\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 14, Prediction Loss: 4.170, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.170\n",
+      "Epoch: 15, Prediction Loss: 4.054, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 4.054\n",
+      "Epoch: 16, Prediction Loss: 3.949, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.949\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 17, Prediction Loss: 3.940, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.940\n",
+      "Epoch: 18, Prediction Loss: 3.918, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.918\n",
+      "Epoch: 19, Prediction Loss: 3.851, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.851\n",
+      "=== Training stage: 2. Train decoder ===\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 0, Prediction Loss: 34.622, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 34.622\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 1, Prediction Loss: 5.329, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 5.329\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 2, Prediction Loss: 3.826, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.826\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 3, Prediction Loss: 3.767, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.767\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 4, Prediction Loss: 3.746, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.746\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 5, Prediction Loss: 3.741, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.741\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 6, Prediction Loss: 3.726, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.726\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 7, Prediction Loss: 3.694, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.694\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 8, Prediction Loss: 3.723, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.723\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 9, Prediction Loss: 3.730, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.730\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 10, Prediction Loss: 3.670, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.670\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 11, Prediction Loss: 3.691, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.691\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 12, Prediction Loss: 3.726, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.726\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 13, Prediction Loss: 3.701, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.701\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 14, Prediction Loss: 3.696, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.696\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 15, Prediction Loss: 3.668, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.668\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 16, Prediction Loss: 3.612, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.612\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 17, Prediction Loss: 3.718, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.718\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 18, Prediction Loss: 3.717, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.717\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 19, Prediction Loss: 3.658, Lambda: 1.000, Treatment BR Loss: 0.000, Loss: 3.658\n"
      ]
     }
    ],
@@ -483,7 +651,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -502,23 +670,23 @@
       "TimeSeries() with data:\n",
       "                 0         1\n",
       "time_idx                    \n",
-      "5         5.993054  4.916986\n",
-      "6         6.041887  4.971257\n",
-      "7         6.043185  4.972708\n",
-      "8         6.043219  4.972745\n",
-      "9         6.043220  4.972746\n",
-      "10        6.043220  4.972746\n",
+      "5         5.966007  4.997900\n",
+      "6         6.120805  5.138746\n",
+      "7         6.138720  5.155504\n",
+      "8         6.140833  5.157483\n",
+      "9         6.141082  5.157716\n",
+      "10        6.141112  5.157744\n",
       "\n",
       "Treatment scenario 1, [0 0 0 0 0 0]\n",
       "TimeSeries() with data:\n",
       "                 0         1\n",
       "time_idx                    \n",
-      "5         6.361402  4.933103\n",
-      "6         6.410686  4.987370\n",
-      "7         6.412068  4.988914\n",
-      "8         6.412107  4.988956\n",
-      "9         6.412108  4.988957\n",
-      "10        6.412108  4.988957\n",
+      "5         6.303439  4.996541\n",
+      "6         6.467038  5.143897\n",
+      "7         6.486365  5.162077\n",
+      "8         6.488708  5.164285\n",
+      "9         6.488993  5.164552\n",
+      "10        6.489028  5.164585\n",
       "\n"
      ]
     }
@@ -587,8 +755,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/usage/tutorial07_pipeline.ipynb b/tutorials/usage/tutorial07_pipeline.ipynb
index 317e0bee..408fd780 100644
--- a/tutorials/usage/tutorial07_pipeline.ipynb
+++ b/tutorials/usage/tutorial07_pipeline.ipynb
@@ -70,7 +70,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -79,7 +79,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -127,8 +127,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:55:13 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n",
-      "2023-10-10 16:55:13 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
+      "2023-11-16 22:18:41 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | hyperimpute.logger:log_and_print:65 | Iteration imputation: select_model_by_column: True, select_model_by_iteration: True\n"
      ]
     },
     {
@@ -256,59 +262,329 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:55:14 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.688586950302124, validation loss: 0.6886438131332397\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:10| train loss: 0.6998531222343445, validation loss: 0.6888602375984192\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:20| train loss: 0.6862186193466187, validation loss: 0.689547061920166\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:30| train loss: 0.6913824677467346, validation loss: 0.6873278617858887\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:40| train loss: 0.6604955792427063, validation loss: 0.689150333404541\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:50| train loss: 0.6746629476547241, validation loss: 0.6878216862678528\n",
-      "2023-10-10 16:55:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:60| train loss: 0.7084729075431824, validation loss: 0.6860194802284241\n",
-      "2023-10-10 16:55:17 | INFO     | tempor.models.ts_model:_train:379 | Epoch:70| train loss: 0.6851019263267517, validation loss: 0.683639645576477\n",
-      "2023-10-10 16:55:17 | INFO     | tempor.models.ts_model:_train:379 | Epoch:80| train loss: 0.6695638298988342, validation loss: 0.6823517084121704\n",
-      "2023-10-10 16:55:17 | INFO     | tempor.models.ts_model:_train:379 | Epoch:90| train loss: 0.6456660032272339, validation loss: 0.6871115565299988\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n",
-      "2023-10-10 16:55:17 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:43 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.688586950302124, validation loss: 0.6886438131332397\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:10| train loss: 0.6998531222343445, validation loss: 0.6888602375984192\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:20| train loss: 0.6862186789512634, validation loss: 0.689547061920166\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:30| train loss: 0.6913825273513794, validation loss: 0.6873278617858887\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:40| train loss: 0.6604955196380615, validation loss: 0.6891502141952515\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:50| train loss: 0.6746630072593689, validation loss: 0.6878217458724976\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:60| train loss: 0.7084730267524719, validation loss: 0.6860195398330688\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:70| train loss: 0.6851019263267517, validation loss: 0.6836397647857666\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:80| train loss: 0.6695641279220581, validation loss: 0.6823516488075256\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:90| train loss: 0.6456663608551025, validation loss: 0.6871115565299988\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > HyperImpute using inner optimization\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 4\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 5\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |   > Imputation iter 6\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | hyperimpute.logger:log_and_print:65 |      >>>> Early stopping on objective diff iteration\n"
      ]
     },
     {
@@ -408,7 +684,7 @@
        "[100 rows x 1 columns]"
       ]
      },
-     "execution_count": 9,
+     "execution_count": null,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -466,8 +742,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/usage/tutorial08_benchmarks.ipynb b/tutorials/usage/tutorial08_benchmarks.ipynb
index 9ef09b88..5d36dadd 100644
--- a/tutorials/usage/tutorial08_benchmarks.ipynb
+++ b/tutorials/usage/tutorial08_benchmarks.ipynb
@@ -47,21 +47,63 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:56:34 | INFO     | tempor.benchmarks.benchmark:benchmark_models:100 | Test case: model_1\n",
-      "2023-10-10 16:56:36 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.68637615442276, validation loss: 0.6680578589439392\n",
-      "2023-10-10 16:56:36 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6887814402580261, validation loss: 0.6939960718154907\n",
-      "2023-10-10 16:56:37 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6906945109367371, validation loss: 0.6966598033905029\n",
-      "2023-10-10 16:56:37 | INFO     | tempor.benchmarks.benchmark:benchmark_models:100 | Test case: model_2\n",
-      "2023-10-10 16:56:41 | INFO     | tempor.models.ts_ode:_train:607 | Epoch:99| train loss: 0.6465951204299927, validation loss: 0.5632617473602295\n",
-      "2023-10-10 16:56:46 | INFO     | tempor.models.ts_ode:_train:607 | Epoch:99| train loss: 0.913261890411377, validation loss: 0.8132617473602295\n",
-      "2023-10-10 16:56:50 | INFO     | tempor.models.ts_ode:_train:607 | Epoch:99| train loss: 0.7132617831230164, validation loss: 0.8132617473602295\n"
+      "2023-11-16 22:18:41 | INFO     | tempor.benchmarks.benchmark:benchmark_models:104 | Test case: model_1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:45 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.68637615442276, validation loss: 0.6680578589439392\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:46 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6887814402580261, validation loss: 0.6939960718154907\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6906945109367371, validation loss: 0.6966598033905029\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:47 | INFO     | tempor.benchmarks.benchmark:benchmark_models:104 | Test case: model_2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:52 | INFO     | tempor.models.ts_ode:_train:617 | Epoch:99| train loss: 0.6465951204299927, validation loss: 0.5632617473602295\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:58 | INFO     | tempor.models.ts_ode:_train:617 | Epoch:99| train loss: 0.913261890411377, validation loss: 0.8132617473602295\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:04 | INFO     | tempor.models.ts_ode:_train:617 | Epoch:99| train loss: 0.7132617831230164, validation loss: 0.8132617473602295\n"
      ]
     },
     {
@@ -98,16 +140,6 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>aucroc</th>\n",
-       "      <td>0.311 +/- 0.083</td>\n",
-       "      <td>0.583 +/- 0.312</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>aucprc</th>\n",
-       "      <td>0.385 +/- 0.028</td>\n",
-       "      <td>0.588 +/- 0.292</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
        "      <th>accuracy</th>\n",
        "      <td>0.602 +/- 0.033</td>\n",
        "      <td>0.519 +/- 0.105</td>\n",
@@ -138,61 +170,71 @@
        "      <td>0.0 +/- 0.0</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>precision_micro</th>\n",
+       "      <th>recall_micro</th>\n",
        "      <td>0.602 +/- 0.033</td>\n",
        "      <td>0.519 +/- 0.105</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>precision_macro</th>\n",
-       "      <td>0.301 +/- 0.016</td>\n",
-       "      <td>0.259 +/- 0.053</td>\n",
+       "      <th>recall_macro</th>\n",
+       "      <td>0.5 +/- 0.0</td>\n",
+       "      <td>0.5 +/- 0.0</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>precision_weighted</th>\n",
-       "      <td>0.363 +/- 0.039</td>\n",
-       "      <td>0.28 +/- 0.104</td>\n",
+       "      <th>recall_weighted</th>\n",
+       "      <td>0.602 +/- 0.033</td>\n",
+       "      <td>0.519 +/- 0.105</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_micro</th>\n",
+       "      <th>precision_micro</th>\n",
        "      <td>0.602 +/- 0.033</td>\n",
        "      <td>0.519 +/- 0.105</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_macro</th>\n",
-       "      <td>0.5 +/- 0.0</td>\n",
-       "      <td>0.5 +/- 0.0</td>\n",
+       "      <th>precision_macro</th>\n",
+       "      <td>0.301 +/- 0.016</td>\n",
+       "      <td>0.259 +/- 0.053</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_weighted</th>\n",
-       "      <td>0.602 +/- 0.033</td>\n",
-       "      <td>0.519 +/- 0.105</td>\n",
+       "      <th>precision_weighted</th>\n",
+       "      <td>0.363 +/- 0.039</td>\n",
+       "      <td>0.28 +/- 0.104</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>mcc</th>\n",
        "      <td>0.0 +/- 0.0</td>\n",
        "      <td>0.0 +/- 0.0</td>\n",
        "    </tr>\n",
+       "    <tr>\n",
+       "      <th>aucprc</th>\n",
+       "      <td>0.385 +/- 0.028</td>\n",
+       "      <td>0.588 +/- 0.292</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>aucroc</th>\n",
+       "      <td>0.311 +/- 0.083</td>\n",
+       "      <td>0.583 +/- 0.312</td>\n",
+       "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "                            model_1          model_2\n",
-       "aucroc              0.311 +/- 0.083  0.583 +/- 0.312\n",
-       "aucprc              0.385 +/- 0.028  0.588 +/- 0.292\n",
        "accuracy            0.602 +/- 0.033  0.519 +/- 0.105\n",
        "f1_score_micro      0.602 +/- 0.033  0.519 +/- 0.105\n",
        "f1_score_macro      0.375 +/- 0.013  0.338 +/- 0.048\n",
        "f1_score_weighted    0.453 +/- 0.04  0.361 +/- 0.116\n",
        "kappa                   0.0 +/- 0.0      0.0 +/- 0.0\n",
        "kappa_quadratic         0.0 +/- 0.0      0.0 +/- 0.0\n",
-       "precision_micro     0.602 +/- 0.033  0.519 +/- 0.105\n",
-       "precision_macro     0.301 +/- 0.016  0.259 +/- 0.053\n",
-       "precision_weighted  0.363 +/- 0.039   0.28 +/- 0.104\n",
        "recall_micro        0.602 +/- 0.033  0.519 +/- 0.105\n",
        "recall_macro            0.5 +/- 0.0      0.5 +/- 0.0\n",
        "recall_weighted     0.602 +/- 0.033  0.519 +/- 0.105\n",
-       "mcc                     0.0 +/- 0.0      0.0 +/- 0.0"
+       "precision_micro     0.602 +/- 0.033  0.519 +/- 0.105\n",
+       "precision_macro     0.301 +/- 0.016  0.259 +/- 0.053\n",
+       "precision_weighted  0.363 +/- 0.039   0.28 +/- 0.104\n",
+       "mcc                     0.0 +/- 0.0      0.0 +/- 0.0\n",
+       "aucprc              0.385 +/- 0.028  0.588 +/- 0.292\n",
+       "aucroc              0.311 +/- 0.083  0.583 +/- 0.312"
       ]
      },
      "metadata": {},
@@ -234,16 +276,6 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>aucroc</th>\n",
-       "      <td>0.311111</td>\n",
-       "      <td>0.083148</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>aucprc</th>\n",
-       "      <td>0.385185</td>\n",
-       "      <td>0.028154</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
        "      <th>accuracy</th>\n",
        "      <td>0.601852</td>\n",
        "      <td>0.032736</td>\n",
@@ -274,6 +306,21 @@
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
+       "      <th>recall_micro</th>\n",
+       "      <td>0.601852</td>\n",
+       "      <td>0.032736</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>recall_macro</th>\n",
+       "      <td>0.500000</td>\n",
+       "      <td>0.000000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>recall_weighted</th>\n",
+       "      <td>0.601852</td>\n",
+       "      <td>0.032736</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
        "      <th>precision_micro</th>\n",
        "      <td>0.601852</td>\n",
        "      <td>0.032736</td>\n",
@@ -289,24 +336,19 @@
        "      <td>0.038647</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_micro</th>\n",
-       "      <td>0.601852</td>\n",
-       "      <td>0.032736</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>recall_macro</th>\n",
-       "      <td>0.500000</td>\n",
+       "      <th>mcc</th>\n",
+       "      <td>0.000000</td>\n",
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_weighted</th>\n",
-       "      <td>0.601852</td>\n",
-       "      <td>0.032736</td>\n",
+       "      <th>aucprc</th>\n",
+       "      <td>0.385185</td>\n",
+       "      <td>0.028154</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>mcc</th>\n",
-       "      <td>0.000000</td>\n",
-       "      <td>0.000000</td>\n",
+       "      <th>aucroc</th>\n",
+       "      <td>0.311111</td>\n",
+       "      <td>0.083148</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -314,21 +356,21 @@
       ],
       "text/plain": [
        "                        mean    stddev\n",
-       "aucroc              0.311111  0.083148\n",
-       "aucprc              0.385185  0.028154\n",
        "accuracy            0.601852  0.032736\n",
        "f1_score_micro      0.601852  0.032736\n",
        "f1_score_macro      0.375458  0.012951\n",
        "f1_score_weighted   0.452788  0.039572\n",
        "kappa               0.000000  0.000000\n",
        "kappa_quadratic     0.000000  0.000000\n",
-       "precision_micro     0.601852  0.032736\n",
-       "precision_macro     0.300926  0.016368\n",
-       "precision_weighted  0.363297  0.038647\n",
        "recall_micro        0.601852  0.032736\n",
        "recall_macro        0.500000  0.000000\n",
        "recall_weighted     0.601852  0.032736\n",
-       "mcc                 0.000000  0.000000"
+       "precision_micro     0.601852  0.032736\n",
+       "precision_macro     0.300926  0.016368\n",
+       "precision_weighted  0.363297  0.038647\n",
+       "mcc                 0.000000  0.000000\n",
+       "aucprc              0.385185  0.028154\n",
+       "aucroc              0.311111  0.083148"
       ]
      },
      "metadata": {},
@@ -368,16 +410,6 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>aucroc</th>\n",
-       "      <td>0.583333</td>\n",
-       "      <td>0.311805</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>aucprc</th>\n",
-       "      <td>0.587731</td>\n",
-       "      <td>0.291568</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
        "      <th>accuracy</th>\n",
        "      <td>0.518519</td>\n",
        "      <td>0.105369</td>\n",
@@ -408,6 +440,21 @@
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
+       "      <th>recall_micro</th>\n",
+       "      <td>0.518519</td>\n",
+       "      <td>0.105369</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>recall_macro</th>\n",
+       "      <td>0.500000</td>\n",
+       "      <td>0.000000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>recall_weighted</th>\n",
+       "      <td>0.518519</td>\n",
+       "      <td>0.105369</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
        "      <th>precision_micro</th>\n",
        "      <td>0.518519</td>\n",
        "      <td>0.105369</td>\n",
@@ -423,24 +470,19 @@
        "      <td>0.104057</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_micro</th>\n",
-       "      <td>0.518519</td>\n",
-       "      <td>0.105369</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>recall_macro</th>\n",
-       "      <td>0.500000</td>\n",
+       "      <th>mcc</th>\n",
+       "      <td>0.000000</td>\n",
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>recall_weighted</th>\n",
-       "      <td>0.518519</td>\n",
-       "      <td>0.105369</td>\n",
+       "      <th>aucprc</th>\n",
+       "      <td>0.587731</td>\n",
+       "      <td>0.291568</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>mcc</th>\n",
-       "      <td>0.000000</td>\n",
-       "      <td>0.000000</td>\n",
+       "      <th>aucroc</th>\n",
+       "      <td>0.583333</td>\n",
+       "      <td>0.311805</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -448,21 +490,21 @@
       ],
       "text/plain": [
        "                        mean    stddev\n",
-       "aucroc              0.583333  0.311805\n",
-       "aucprc              0.587731  0.291568\n",
        "accuracy            0.518519  0.105369\n",
        "f1_score_micro      0.518519  0.105369\n",
        "f1_score_macro      0.338162  0.047609\n",
        "f1_score_weighted   0.360713  0.115623\n",
        "kappa               0.000000  0.000000\n",
        "kappa_quadratic     0.000000  0.000000\n",
-       "precision_micro     0.518519  0.105369\n",
-       "precision_macro     0.259259  0.052684\n",
-       "precision_weighted  0.279964  0.104057\n",
        "recall_micro        0.518519  0.105369\n",
        "recall_macro        0.500000  0.000000\n",
        "recall_weighted     0.518519  0.105369\n",
-       "mcc                 0.000000  0.000000"
+       "precision_micro     0.518519  0.105369\n",
+       "precision_macro     0.259259  0.052684\n",
+       "precision_weighted  0.279964  0.104057\n",
+       "mcc                 0.000000  0.000000\n",
+       "aucprc              0.587731  0.291568\n",
+       "aucroc              0.583333  0.311805"
       ]
      },
      "metadata": {},
@@ -552,8 +594,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
diff --git a/tutorials/usage/tutorial09_automl.ipynb b/tutorials/usage/tutorial09_automl.ipynb
index 752b7946..f7c52477 100644
--- a/tutorials/usage/tutorial09_automl.ipynb
+++ b/tutorials/usage/tutorial09_automl.ipynb
@@ -65,17 +65,35 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator cde_classifier.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator ode_classifier.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator nn_classifier.\n"
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator cde_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator ode_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator nn_classifier.\n"
      ]
     }
    ],
@@ -124,17 +142,35 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator cde_classifier.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator ode_classifier.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator nn_classifier.\n"
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator cde_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator ode_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator nn_classifier.\n"
      ]
     }
    ],
@@ -182,89 +218,317 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator 'cde_classifier' 1/3.\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:57:26 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator 'cde_classifier' 1/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:41 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from CDEClassifier:\n",
-      "{'n_iter': 2, 'n_temporal_units_hidden': 13, 'lr': 0.01}\n",
-      "2023-10-10 16:57:31 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_temporal_units_hidden': 13, 'lr': 0.01}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:48 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from CDEClassifier:\n",
-      "{'n_iter': 2, 'n_temporal_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:34 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_temporal_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:53 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from CDEClassifier:\n",
-      "{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}\n",
-      "2023-10-10 16:57:38 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator 'ode_classifier' 2/3.\n",
-      "2023-10-10 16:57:38 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:57:38 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:58 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator 'ode_classifier' 2/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:58 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:18:58 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from ODEClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 13, 'lr': 0.01}\n",
-      "2023-10-10 16:57:41 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 13, 'lr': 0.01}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:01 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from ODEClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:44 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:05 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from ODEClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 20, 'lr': 0.001}\n",
-      "2023-10-10 16:57:48 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator 'nn_classifier' 3/3.\n",
-      "2023-10-10 16:57:48 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:57:48 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 20, 'lr': 0.001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:08 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator 'nn_classifier' 3/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:08 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:08 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from NeuralNetClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 13, 'lr': 0.01}\n",
-      "2023-10-10 16:57:48 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.6882655620574951\n",
-      "2023-10-10 16:57:49 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6865323185920715\n",
-      "2023-10-10 16:57:49 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6861070990562439\n",
-      "2023-10-10 16:57:50 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6825112104415894\n",
-      "2023-10-10 16:57:50 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6819757223129272\n",
-      "2023-10-10 16:57:50 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 13, 'lr': 0.01}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:09 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.6882655620574951\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:09 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6865323185920715\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:10 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6861070990562439\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:10 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6825112104415894\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:11 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6819757223129272\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:11 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from NeuralNetClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:51 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.691270649433136\n",
-      "2023-10-10 16:57:52 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6910805106163025\n",
-      "2023-10-10 16:57:52 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6900344491004944\n",
-      "2023-10-10 16:57:53 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6906693577766418\n",
-      "2023-10-10 16:57:53 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6910874843597412\n",
-      "2023-10-10 16:57:53 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:12 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.691270649433136\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:12 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6910805106163025\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:13 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6900344491004944\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:13 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6906693577766418\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:14 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6910874843597412\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:14 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from NeuralNetClassifier:\n",
-      "{'n_iter': 2, 'n_units_hidden': 20, 'lr': 0.001}\n",
-      "2023-10-10 16:57:54 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.6896671056747437\n",
-      "2023-10-10 16:57:54 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6886401772499084\n",
-      "2023-10-10 16:57:55 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6884845495223999\n",
-      "2023-10-10 16:57:55 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6887120604515076\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6893669962882996\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "{'n_iter': 2, 'n_units_hidden': 20, 'lr': 0.001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:14 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919505596160889, validation loss: 0.6896671056747437\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:15 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.691156268119812, validation loss: 0.6886401772499084\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:15 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.692247211933136, validation loss: 0.6884845495223999\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6928325891494751, validation loss: 0.6887120604515076\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6919053792953491, validation loss: 0.6893669962882996\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for cde_classifier scores:\n",
-      "[0.47310606060606064, 0.48396464646464643, 0.4840277777777778].\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "[0.46300505050505053, 0.48396464646464643, 0.4901515151515151].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for ode_classifier scores:\n",
-      "[0.47992424242424236, 0.48409090909090907, 0.484090909090909].\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "[0.47992424242424236, 0.48409090909090907, 0.484090909090909].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for nn_classifier scores:\n",
-      "[0.37064393939393936, 0.5370580808080809, 0.4579545454545454].\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:460 | \n",
+      "[0.37064393939393936, 0.5370580808080809, 0.4579545454545454].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:436 | \n",
       "All estimator definitions searched:\n",
-      "['cde_classifier', 'ode_classifier', 'nn_classifier']\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:461 | \n",
+      "['cde_classifier', 'ode_classifier', 'nn_classifier']\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:437 | \n",
       "Best scores for each estimator searched:\n",
-      "[0.4840277777777778, 0.48409090909090907, 0.5370580808080809]\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:462 | \n",
+      "[0.4901515151515151, 0.48409090909090907, 0.5370580808080809]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:search:438 | \n",
       "Best hyperparameters for each estimator searched:\n",
-      "[{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}, {'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}, {'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}]\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:587 | \n",
+      "[{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}, {'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}, {'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:565 | \n",
       "Selected score 0.5370580808080809 for nn_classifier with hyperparameters:\n",
-      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:587 | \n",
-      "Selected score 0.48409090909090907 for ode_classifier with hyperparameters:\n",
-      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:587 | \n",
-      "Selected score 0.4840277777777778 for cde_classifier with hyperparameters:\n",
+      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:565 | \n",
+      "Selected score 0.4901515151515151 for cde_classifier with hyperparameters:\n",
       "{'n_iter': 2, 'n_temporal_units_hidden': 20, 'lr': 0.001}\n"
      ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:565 | \n",
+      "Selected score 0.48409090909090907 for ode_classifier with hyperparameters:\n",
+      "{'n_iter': 2, 'n_units_hidden': 11, 'lr': 0.0001}\n"
+     ]
     }
    ],
    "source": [
@@ -275,7 +539,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -344,19 +608,19 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ODEClassifier</span><span style=\"font-weight: bold\">(</span>\n",
-       "    <span style=\"color: #808000; text-decoration-color: #808000\">name</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'ode_classifier'</span>,\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">CDEClassifier</span><span style=\"font-weight: bold\">(</span>\n",
+       "    <span style=\"color: #808000; text-decoration-color: #808000\">name</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'cde_classifier'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">category</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'prediction.one_off.classification'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">plugin_type</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'method'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">params</span>=<span style=\"font-weight: bold\">{</span>\n",
-       "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_units_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_units_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">100</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_layers_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'nonlin'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'relu'</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'dropout'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'atol'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'rtol'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'interpolation'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'cubic'</span>,\n",
-       "        <span style=\"color: #008000; text-decoration-color: #008000\">'lr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0001</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'lr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.001</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'weight_decay'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.001</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_iter'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'batch_size'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">500</span>,\n",
@@ -372,19 +636,19 @@
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[1;35mODEClassifier\u001b[0m\u001b[1m(\u001b[0m\n",
-       "    \u001b[33mname\u001b[0m=\u001b[32m'ode_classifier'\u001b[0m,\n",
+       "\u001b[1;35mCDEClassifier\u001b[0m\u001b[1m(\u001b[0m\n",
+       "    \u001b[33mname\u001b[0m=\u001b[32m'cde_classifier'\u001b[0m,\n",
        "    \u001b[33mcategory\u001b[0m=\u001b[32m'prediction.one_off.classification'\u001b[0m,\n",
        "    \u001b[33mplugin_type\u001b[0m=\u001b[32m'method'\u001b[0m,\n",
        "    \u001b[33mparams\u001b[0m=\u001b[1m{\u001b[0m\n",
-       "        \u001b[32m'n_units_hidden'\u001b[0m: \u001b[1;36m11\u001b[0m,\n",
+       "        \u001b[32m'n_units_hidden'\u001b[0m: \u001b[1;36m100\u001b[0m,\n",
        "        \u001b[32m'n_layers_hidden'\u001b[0m: \u001b[1;36m1\u001b[0m,\n",
        "        \u001b[32m'nonlin'\u001b[0m: \u001b[32m'relu'\u001b[0m,\n",
        "        \u001b[32m'dropout'\u001b[0m: \u001b[1;36m0\u001b[0m,\n",
        "        \u001b[32m'atol'\u001b[0m: \u001b[1;36m0.01\u001b[0m,\n",
        "        \u001b[32m'rtol'\u001b[0m: \u001b[1;36m0.01\u001b[0m,\n",
        "        \u001b[32m'interpolation'\u001b[0m: \u001b[32m'cubic'\u001b[0m,\n",
-       "        \u001b[32m'lr'\u001b[0m: \u001b[1;36m0.0001\u001b[0m,\n",
+       "        \u001b[32m'lr'\u001b[0m: \u001b[1;36m0.001\u001b[0m,\n",
        "        \u001b[32m'weight_decay'\u001b[0m: \u001b[1;36m0.001\u001b[0m,\n",
        "        \u001b[32m'n_iter'\u001b[0m: \u001b[1;36m2\u001b[0m,\n",
        "        \u001b[32m'batch_size'\u001b[0m: \u001b[1;36m500\u001b[0m,\n",
@@ -405,19 +669,19 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">CDEClassifier</span><span style=\"font-weight: bold\">(</span>\n",
-       "    <span style=\"color: #808000; text-decoration-color: #808000\">name</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'cde_classifier'</span>,\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ODEClassifier</span><span style=\"font-weight: bold\">(</span>\n",
+       "    <span style=\"color: #808000; text-decoration-color: #808000\">name</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'ode_classifier'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">category</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'prediction.one_off.classification'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">plugin_type</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'method'</span>,\n",
        "    <span style=\"color: #808000; text-decoration-color: #808000\">params</span>=<span style=\"font-weight: bold\">{</span>\n",
-       "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_units_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">100</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_units_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_layers_hidden'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'nonlin'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'relu'</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'dropout'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'atol'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'rtol'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'interpolation'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'cubic'</span>,\n",
-       "        <span style=\"color: #008000; text-decoration-color: #008000\">'lr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.001</span>,\n",
+       "        <span style=\"color: #008000; text-decoration-color: #008000\">'lr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0001</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'weight_decay'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.001</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'n_iter'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,\n",
        "        <span style=\"color: #008000; text-decoration-color: #008000\">'batch_size'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">500</span>,\n",
@@ -433,19 +697,19 @@
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[1;35mCDEClassifier\u001b[0m\u001b[1m(\u001b[0m\n",
-       "    \u001b[33mname\u001b[0m=\u001b[32m'cde_classifier'\u001b[0m,\n",
+       "\u001b[1;35mODEClassifier\u001b[0m\u001b[1m(\u001b[0m\n",
+       "    \u001b[33mname\u001b[0m=\u001b[32m'ode_classifier'\u001b[0m,\n",
        "    \u001b[33mcategory\u001b[0m=\u001b[32m'prediction.one_off.classification'\u001b[0m,\n",
        "    \u001b[33mplugin_type\u001b[0m=\u001b[32m'method'\u001b[0m,\n",
        "    \u001b[33mparams\u001b[0m=\u001b[1m{\u001b[0m\n",
-       "        \u001b[32m'n_units_hidden'\u001b[0m: \u001b[1;36m100\u001b[0m,\n",
+       "        \u001b[32m'n_units_hidden'\u001b[0m: \u001b[1;36m11\u001b[0m,\n",
        "        \u001b[32m'n_layers_hidden'\u001b[0m: \u001b[1;36m1\u001b[0m,\n",
        "        \u001b[32m'nonlin'\u001b[0m: \u001b[32m'relu'\u001b[0m,\n",
        "        \u001b[32m'dropout'\u001b[0m: \u001b[1;36m0\u001b[0m,\n",
        "        \u001b[32m'atol'\u001b[0m: \u001b[1;36m0.01\u001b[0m,\n",
        "        \u001b[32m'rtol'\u001b[0m: \u001b[1;36m0.01\u001b[0m,\n",
        "        \u001b[32m'interpolation'\u001b[0m: \u001b[32m'cubic'\u001b[0m,\n",
-       "        \u001b[32m'lr'\u001b[0m: \u001b[1;36m0.001\u001b[0m,\n",
+       "        \u001b[32m'lr'\u001b[0m: \u001b[1;36m0.0001\u001b[0m,\n",
        "        \u001b[32m'weight_decay'\u001b[0m: \u001b[1;36m0.001\u001b[0m,\n",
        "        \u001b[32m'n_iter'\u001b[0m: \u001b[1;36m2\u001b[0m,\n",
        "        \u001b[32m'batch_size'\u001b[0m: \u001b[1;36m500\u001b[0m,\n",
@@ -497,17 +761,35 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with cde_classifier>.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with ode_classifier>.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with nn_classifier>.\n"
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with cde_classifier>.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with ode_classifier>.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:16 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with nn_classifier>.\n"
      ]
     }
    ],
@@ -546,7 +828,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -584,17 +866,35 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with cde_classifier>.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with ode_classifier>.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:_init_estimator:753 | Creating estimator <Pipeline with nn_classifier>.\n"
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with cde_classifier>.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with ode_classifier>.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:_init_estimator:733 | Creating estimator <Pipeline with nn_classifier>.\n"
      ]
     }
    ],
@@ -649,87 +949,315 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator '<Pipeline with cde_classifier>' 1/3.\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:57:56 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator '<Pipeline with cde_classifier>' 1/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:17 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.cde_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 17, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:14 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 17, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:32 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_minmax_scaler->prediction.one_off.classification.cde_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 14, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:20 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 14, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:38 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_standard_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.cde_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 6, 'lr': 0.0001}}}\n",
-      "2023-10-10 16:58:26 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator '<Pipeline with ode_classifier>' 2/3.\n",
-      "2023-10-10 16:58:26 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:58:26 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'cde_classifier': {'n_iter': 2, 'n_temporal_units_hidden': 6, 'lr': 0.0001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:45 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator '<Pipeline with ode_classifier>' 2/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:45 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:19:45 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.ode_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 17, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:45 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 17, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:00 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_minmax_scaler->prediction.one_off.classification.ode_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 14, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:50 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 14, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:06 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_standard_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.ode_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 6, 'lr': 0.0001}}}\n",
-      "2023-10-10 16:58:56 | INFO     | tempor.automl.seeker:search:424 | Running  search for estimator '<Pipeline with nn_classifier>' 3/3.\n",
-      "2023-10-10 16:58:56 | INFO     | tempor.automl.tuner:tune:205 | Baseline score computation skipped\n",
-      "2023-10-10 16:58:56 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'ode_classifier': {'n_iter': 2, 'n_units_hidden': 6, 'lr': 0.0001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:12 | INFO     | tempor.automl.seeker:search:400 | Running  search for estimator '<Pipeline with nn_classifier>' 3/3.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:12 | INFO     | tempor.automl.tuner:tune:228 | Baseline score computation skipped\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:12 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.nn_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 17, 'lr': 0.001}}}\n",
-      "2023-10-10 16:58:58 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6918608546257019, validation loss: 0.6910383105278015\n",
-      "2023-10-10 16:59:01 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6909429430961609, validation loss: 0.6872987151145935\n",
-      "2023-10-10 16:59:04 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.692150354385376, validation loss: 0.6879752278327942\n",
-      "2023-10-10 16:59:07 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6926363706588745, validation loss: 0.6881052851676941\n",
-      "2023-10-10 16:59:10 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6915659308433533, validation loss: 0.6888220310211182\n",
-      "2023-10-10 16:59:11 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'softimpute'}, 'static_minmax_scaler': {'clip': True}, 'ts_standard_scaler': {}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 17, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:14 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6918608546257019, validation loss: 0.6910383105278015\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:17 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6909429430961609, validation loss: 0.6872987151145935\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:20 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.692150354385376, validation loss: 0.6879752278327942\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:23 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6926363706588745, validation loss: 0.6881052851676941\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:26 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6915659308433533, validation loss: 0.6888220310211182\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:27 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_minmax_scaler->preprocessing.scaling.temporal.ts_minmax_scaler->prediction.one_off.classification.nn_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 14, 'lr': 0.001}}}\n",
-      "2023-10-10 16:59:12 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6920319199562073, validation loss: 0.6892786622047424\n",
-      "2023-10-10 16:59:13 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6913070678710938, validation loss: 0.6893746852874756\n",
-      "2023-10-10 16:59:14 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6924894452095032, validation loss: 0.6886722445487976\n",
-      "2023-10-10 16:59:15 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6930248141288757, validation loss: 0.6894596815109253\n",
-      "2023-10-10 16:59:16 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6921922564506531, validation loss: 0.6893032789230347\n",
-      "2023-10-10 16:59:16 | INFO     | tempor.automl.tuner:objective:227 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'mean'}, 'static_minmax_scaler': {'clip': False}, 'ts_minmax_scaler': {'clip': False}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 14, 'lr': 0.001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:27 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6920319199562073, validation loss: 0.6892786622047424\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:28 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6913070678710938, validation loss: 0.6893746852874756\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:30 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6924894452095032, validation loss: 0.6886722445487976\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:31 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6930248141288757, validation loss: 0.6894596815109253\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:32 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6921922564506531, validation loss: 0.6893032789230347\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:32 | INFO     | tempor.automl.tuner:objective:248 | \n",
       "Hyperparameters sampled from preprocessing.imputation.static.static_tabular_imputer->preprocessing.scaling.static.static_standard_scaler->preprocessing.scaling.temporal.ts_standard_scaler->prediction.one_off.classification.nn_classifier:\n",
-      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 6, 'lr': 0.0001}}}\n",
-      "2023-10-10 16:59:17 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6914154291152954, validation loss: 0.6911771297454834\n",
-      "2023-10-10 16:59:18 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.690178394317627, validation loss: 0.6934390664100647\n",
-      "2023-10-10 16:59:19 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6933540105819702, validation loss: 0.6896734833717346\n",
-      "2023-10-10 16:59:20 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.694333553314209, validation loss: 0.690070390701294\n",
-      "2023-10-10 16:59:21 | INFO     | tempor.models.ts_model:_train:379 | Epoch:0| train loss: 0.6921719312667847, validation loss: 0.6909489631652832\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "{'plugin_params': {'static_tabular_imputer': {'imputer': 'most_frequent'}, 'static_standard_scaler': {}, 'ts_standard_scaler': {}, 'nn_classifier': {'n_iter': 2, 'n_units_hidden': 6, 'lr': 0.0001}}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:33 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6914154291152954, validation loss: 0.6911771297454834\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:34 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.690178394317627, validation loss: 0.6934390664100647\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:35 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6933540105819702, validation loss: 0.6896734833717346\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:36 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.694333553314209, validation loss: 0.690070390701294\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.models.ts_model:_train:388 | Epoch:0| train loss: 0.6921719312667847, validation loss: 0.6909489631652832\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for <Pipeline with cde_classifier> scores:\n",
-      "[0.4716540404040404, 0.47575757575757577, 0.4657828282828283].\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "[0.4736742424242424, 0.4737373737373737, 0.4657828282828283].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for <Pipeline with ode_classifier> scores:\n",
-      "[0.47575757575757577, 0.4779040404040404, 0.4779671717171716].\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:459 | \n",
+      "[0.47575757575757577, 0.4779040404040404, 0.4779671717171716].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:435 | \n",
       "Evaluation for <Pipeline with nn_classifier> scores:\n",
-      "[0.48478535353535346, 0.4361742424242424, 0.5071338383838383].\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:460 | \n",
+      "[0.48478535353535346, 0.4361742424242424, 0.5071338383838383].\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:436 | \n",
       "All estimator definitions searched:\n",
-      "['<Pipeline with cde_classifier>', '<Pipeline with ode_classifier>', '<Pipeline with nn_classifier>']\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:461 | \n",
+      "['<Pipeline with cde_classifier>', '<Pipeline with ode_classifier>', '<Pipeline with nn_classifier>']\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:437 | \n",
       "Best scores for each estimator searched:\n",
-      "[0.47575757575757577, 0.4779671717171716, 0.5071338383838383]\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:search:462 | \n",
+      "[0.4737373737373737, 0.4779671717171716, 0.5071338383838383]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:search:438 | \n",
       "Best hyperparameters for each estimator searched:\n",
-      "[{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'mean', '<candidates>(preprocessing.scaling.static)': 'static_minmax_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_minmax_scaler', '[ts_minmax_scaler](clip)': False, '[cde_classifier](n_iter)': 2, '[cde_classifier](n_temporal_units_hidden)': 14, '[cde_classifier](lr)': 0.001}, {'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[ode_classifier](n_iter)': 2, '[ode_classifier](n_units_hidden)': 6, '[ode_classifier](lr)': 0.0001}, {'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[nn_classifier](n_iter)': 2, '[nn_classifier](n_units_hidden)': 6, '[nn_classifier](lr)': 0.0001}]\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:768 | \n",
+      "[{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'mean', '<candidates>(preprocessing.scaling.static)': 'static_minmax_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_minmax_scaler', '[ts_minmax_scaler](clip)': False, '[cde_classifier](n_iter)': 2, '[cde_classifier](n_temporal_units_hidden)': 14, '[cde_classifier](lr)': 0.001}, {'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[ode_classifier](n_iter)': 2, '[ode_classifier](n_units_hidden)': 6, '[ode_classifier](lr)': 0.0001}, {'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[nn_classifier](n_iter)': 2, '[nn_classifier](n_units_hidden)': 6, '[nn_classifier](lr)': 0.0001}]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:748 | \n",
       "Selected score 0.5071338383838383 for <Pipeline with nn_classifier> with hyperparameters:\n",
-      "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[nn_classifier](n_iter)': 2, '[nn_classifier](n_units_hidden)': 6, '[nn_classifier](lr)': 0.0001}\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:768 | \n",
+      "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[nn_classifier](n_iter)': 2, '[nn_classifier](n_units_hidden)': 6, '[nn_classifier](lr)': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:748 | \n",
       "Selected score 0.4779671717171716 for <Pipeline with ode_classifier> with hyperparameters:\n",
-      "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[ode_classifier](n_iter)': 2, '[ode_classifier](n_units_hidden)': 6, '[ode_classifier](lr)': 0.0001}\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:768 | \n",
-      "Selected score 0.47575757575757577 for <Pipeline with cde_classifier> with hyperparameters:\n",
+      "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'most_frequent', '<candidates>(preprocessing.scaling.static)': 'static_standard_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_standard_scaler', '[ts_minmax_scaler](clip)': False, '[ode_classifier](n_iter)': 2, '[ode_classifier](n_units_hidden)': 6, '[ode_classifier](lr)': 0.0001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_create_estimator_with_hps:748 | \n",
+      "Selected score 0.4737373737373737 for <Pipeline with cde_classifier> with hyperparameters:\n",
       "{'<candidates>(preprocessing.imputation.static)': 'static_tabular_imputer', '[static_tabular_imputer](imputer)': 'mean', '<candidates>(preprocessing.scaling.static)': 'static_minmax_scaler', '[static_minmax_scaler](clip)': False, '<candidates>(preprocessing.scaling.temporal)': 'ts_minmax_scaler', '[ts_minmax_scaler](clip)': False, '[cde_classifier](n_iter)': 2, '[cde_classifier](n_temporal_units_hidden)': 14, '[cde_classifier](lr)': 0.001}\n"
      ]
     }
@@ -740,7 +1268,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -990,17 +1518,35 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_set_up_tuners:365 | Setting up estimators and tuners for study my_automl_study.\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator cde_classifier.\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator ode_classifier.\n",
-      "2023-10-10 16:59:22 | INFO     | tempor.automl.seeker:_init_estimator:579 | Creating estimator nn_classifier.\n"
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_set_up_tuners:341 | Setting up estimators and tuners for study my_automl_study.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator cde_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator ode_classifier.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-11-16 22:20:37 | INFO     | tempor.automl.seeker:_init_estimator:557 | Creating estimator nn_classifier.\n"
      ]
     }
    ],
@@ -1092,8 +1638,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2


Test patch:
diff --git a/tests/automl/test_seeker.py b/tests/automl/test_seeker.py
index acdfbe88..9dcc68f7 100644
--- a/tests/automl/test_seeker.py
+++ b/tests/automl/test_seeker.py
@@ -5,11 +5,11 @@
 
 import pytest
 
+from tempor import plugin_loader
 from tempor.automl._types import AutoMLCompatibleEstimator
 from tempor.automl.seeker import (
     DEFAULT_STATIC_SCALERS,
     DEFAULT_TEMPORAL_SCALERS,
-    METRIC_DIRECTION_MAP,
     BaseSeeker,
     MethodSeeker,
     PipelineSeeker,
@@ -83,9 +83,12 @@ def patched_evaluate(*args, **kwargs):
         np.random.seed(hash_)
         # --- --- ---
 
+        metric_plugin_full_names = plugin_loader.list_full_names(plugin_type="metric")
+        metric_names = [x.split(".")[-1] for x in metric_plugin_full_names]
+        num_metrics = len(metric_names)
         return pd.DataFrame(
-            data=np.random.rand(len(METRIC_DIRECTION_MAP), 1),
-            index=list(METRIC_DIRECTION_MAP.keys()),
+            data=np.random.rand(num_metrics, 1),
+            index=metric_names,
             columns=["mean"],
         )
 
diff --git a/tests/automl/test_tuner.py b/tests/automl/test_tuner.py
index 4f0f56cd..9d25fd69 100644
--- a/tests/automl/test_tuner.py
+++ b/tests/automl/test_tuner.py
@@ -10,11 +10,12 @@
 from packaging.version import Version
 
 from tempor import plugin_loader
-from tempor.automl import OptimDirection, tuner
+from tempor.automl import tuner
 from tempor.benchmarks import evaluation
 from tempor.data.dataset import PredictiveDataset, TimeToEventAnalysisDataset
 from tempor.methods.core._base_predictor import BasePredictor
 from tempor.methods.core.params import IntegerParams
+from tempor.metrics import metric_typing
 
 # To ignore warnings in parametrization:
 warnings.filterwarnings("ignore", category=optuna.exceptions.ExperimentalWarning)
@@ -294,7 +295,7 @@ def func(
         pruner: Any,
         evaluation_case: str,
         metric: str,
-        direction: OptimDirection,
+        direction: metric_typing.MetricDirection,
         compute_baseline_score: bool,
     ):
         # A general test function for testing `tuner.OptunaTuner`.
@@ -345,7 +346,7 @@ def func(
         pruner: Any,
         evaluation_case: str,
         metric: str,
-        direction: OptimDirection,
+        direction: metric_typing.MetricDirection,
     ):
         # A general test function for testing `tuner.OptunaTuner` - pipeline selector case.
 
@@ -548,7 +549,7 @@ def test_tune_vary_metrics(
             data: str,
             plugin: str,
             metric: str,
-            direction: OptimDirection,
+            direction: metric_typing.MetricDirection,
             helper_test_optuna_tuner: Callable,
         ):
             helper_test_optuna_tuner(
@@ -605,7 +606,12 @@ def test_tune_vary_pruners(self, data: str, plugin: str, pruner: Any, helper_tes
         @pytest.mark.parametrize("plugin", SETTINGS["prediction.one_off.regression"]["TEST_WITH_PLUGINS"])
         @pytest.mark.parametrize("metric,direction", SETTINGS["prediction.one_off.regression"]["METRICS"])
         def test_tune_vary_metrics(
-            self, data: str, plugin: str, metric: str, direction: OptimDirection, helper_test_optuna_tuner: Callable
+            self,
+            data: str,
+            plugin: str,
+            metric: str,
+            direction: metric_typing.MetricDirection,
+            helper_test_optuna_tuner: Callable,
         ):
             helper_test_optuna_tuner(
                 data=data,
@@ -664,7 +670,12 @@ def test_tune_vary_pruners(self, data: str, plugin: str, pruner: Any, helper_tes
         @pytest.mark.parametrize("plugin", SETTINGS["time_to_event"]["TEST_WITH_PLUGINS"])
         @pytest.mark.parametrize("metric,direction", SETTINGS["time_to_event"]["METRICS"])
         def test_tune_vary_metrics(
-            self, data: str, plugin: str, metric: str, direction: OptimDirection, helper_test_optuna_tuner: Callable
+            self,
+            data: str,
+            plugin: str,
+            metric: str,
+            direction: metric_typing.MetricDirection,
+            helper_test_optuna_tuner: Callable,
         ):
             helper_test_optuna_tuner(
                 data=data,
diff --git a/tests/benchmarks/test_benchmark.py b/tests/benchmarks/test_benchmark.py
index e42a3375..47d7c24f 100644
--- a/tests/benchmarks/test_benchmark.py
+++ b/tests/benchmarks/test_benchmark.py
@@ -6,9 +6,9 @@
 from tempor import plugin_loader
 from tempor.benchmarks import (
     benchmark_models,
-    classifier_supported_metrics,
-    regression_supported_metrics,
-    time_to_event_supported_metrics,
+    builtin_metrics_prediction_oneoff_classification,
+    builtin_metrics_prediction_oneoff_regression,
+    builtin_metrics_time_to_event,
     visualize_benchmark,
 )
 from tempor.methods.pipeline import pipeline
@@ -60,7 +60,7 @@ def test_classifier_benchmark(data: str, request: pytest.FixtureRequest) -> None
         assert testcase in aggr_score.columns
         assert testcase in per_test_score
 
-    for metric in classifier_supported_metrics:
+    for metric in builtin_metrics_prediction_oneoff_classification:
         assert metric in aggr_score.index
 
         for testcase, _ in testcases:
@@ -105,7 +105,7 @@ def test_regressor_benchmark(data: str, request: pytest.FixtureRequest) -> None:
         assert testcase in aggr_score.columns
         assert testcase in per_test_score
 
-    for metric in regression_supported_metrics:
+    for metric in builtin_metrics_prediction_oneoff_regression:
         assert metric in aggr_score.index
 
         for testcase, _ in testcases:
@@ -149,7 +149,7 @@ def test_time_to_event_benchmark(
         assert testcase in aggr_score.columns
         assert testcase in per_test_score
 
-    for metric in time_to_event_supported_metrics:
+    for metric in builtin_metrics_time_to_event:
         assert metric in aggr_score.index
 
         for testcase, _ in testcases:
diff --git a/tests/benchmarks/test_benchmarks_utils.py b/tests/benchmarks/test_benchmarks_utils.py
index a942ed15..add5e32e 100644
--- a/tests/benchmarks/test_benchmarks_utils.py
+++ b/tests/benchmarks/test_benchmarks_utils.py
@@ -1,41 +1,18 @@
-from unittest.mock import Mock
-
 import numpy as np
-import pytest
 
 from tempor.benchmarks import utils
 
 
-def test_get_y_pred_proba_hlpr_edge():
-    mock = Mock()
-    out = utils.get_y_pred_proba_hlpr(
-        y_pred_proba=mock,
-        nclasses=3,
-    )
-    assert out == mock
-
-    mock = Mock(shape=[1])
-    out = utils.get_y_pred_proba_hlpr(
-        y_pred_proba=mock,
-        nclasses=2,
-    )
-    assert out == mock
-
-    mock = Mock(shape=[1, 4])
-    out = utils.get_y_pred_proba_hlpr(
-        y_pred_proba=mock,
-        nclasses=2,
-    )
-    assert out == mock
-
+def test_generate_score():
+    out = utils.generate_score(metric=np.asarray([1.0, 2.0, 3.0]))
+    assert len(out) == 2
 
-def test_evaluate_auc_multiclass_nans():
-    with pytest.raises(ValueError, match=".*nan.*"):
-        utils.evaluate_auc_multiclass(y_test=np.asarray([1, 1, 1]), y_pred_proba=np.asarray([0.9, np.nan, 0.9]))
+    mean, ci = out
+    np.testing.assert_equal(mean, 2.0)
+    np.testing.assert_almost_equal(ci, 0.9239, 3)
 
 
-def test_evaluate_auc_multiclass():
-    out = utils.evaluate_auc_multiclass(
-        y_test=np.asarray([0, 2, 1]), y_pred_proba=np.asarray([[0.8, 0.1, 0.1], [0.1, 0.1, 0.8], [0.1, 0.8, 0.1]])
-    )
-    assert len(out) == 2
+def test_print_score():
+    out = utils.print_score(score=(1, 2))
+    assert isinstance(out, str)
+    assert out == "1 +/- 2"
diff --git a/tests/benchmarks/test_evaluation.py b/tests/benchmarks/test_evaluation.py
index 594290aa..c0bb471c 100644
--- a/tests/benchmarks/test_evaluation.py
+++ b/tests/benchmarks/test_evaluation.py
@@ -5,15 +5,14 @@
 
 from tempor import plugin_loader
 from tempor.benchmarks import (
-    classifier_supported_metrics,
+    builtin_metrics_prediction_oneoff_classification,
+    builtin_metrics_prediction_oneoff_regression,
+    builtin_metrics_time_to_event,
     evaluate_prediction_oneoff_classifier,
     evaluate_prediction_oneoff_regressor,
     evaluate_time_to_event,
     output_metrics,
-    regression_supported_metrics,
-    time_to_event_supported_metrics,
 )
-from tempor.benchmarks.evaluation import ClassifierMetrics
 from tempor.methods.pipeline import pipeline
 
 N_ITER = 5
@@ -27,12 +26,6 @@
 PREDICTOR_TIME_TO_EVENT = "time_to_event.dynamic_deephit"
 
 
-def test_classifier_metrics_score_proba_input_validation():
-    clf_metrics = ClassifierMetrics()
-    with pytest.raises(ValueError, match=".*.input*"):
-        clf_metrics.score_proba(None, None)  # type: ignore
-
-
 @pytest.mark.parametrize("data", TEST_ON_DATASETS_CLASSIFIER)
 @pytest.mark.parametrize(
     "model_template",
@@ -62,7 +55,7 @@ def test_evaluate_prediction_oneoff_classifier(
     for out_metric in output_metrics:
         assert out_metric in scores
 
-    for metric in classifier_supported_metrics:
+    for metric in builtin_metrics_prediction_oneoff_classification:
         assert metric in scores.index
 
 
@@ -161,7 +154,7 @@ def test_evaluate_prediction_oneoff_regressor(
     for out_metric in output_metrics:
         assert out_metric in scores
 
-    for metric in regression_supported_metrics:
+    for metric in builtin_metrics_prediction_oneoff_regression:
         assert metric in scores.index
 
 
@@ -241,7 +234,7 @@ def test_time_to_event_evaluation(
     for out_metric in output_metrics:
         assert out_metric in scores
 
-    for metric in time_to_event_supported_metrics:
+    for metric in builtin_metrics_time_to_event:
         assert metric in scores.index
 
 
@@ -304,16 +297,14 @@ def raise_(*args, **kwargs):
 def test_compute_time_to_event_metric():
     import pandas as pd
 
-    from tempor.benchmarks.evaluation import _compute_time_to_event_metric
+    from tempor.benchmarks.evaluation import _prep_data_for_time_to_event_metric
 
-    metric_func = Mock()
     mock_pred = Mock(num_timesteps_equal=Mock(return_value=True), num_features=1)
     mock_test_data = Mock(predictive=Mock(targets=Mock(num_features=1)))
     mock_train_data = Mock(predictive=Mock(targets=Mock(num_features=1)))
 
     with pytest.raises(ValueError, match=".*horizon.*"):
-        _compute_time_to_event_metric(
-            metric_func=metric_func,
+        _prep_data_for_time_to_event_metric(
             train_data=mock_train_data,
             test_data=mock_test_data,
             horizons=pd.to_datetime(["2000-01-01", "2000-01-02"]),  # type: ignore
@@ -321,8 +312,7 @@ def test_compute_time_to_event_metric():
         )
 
     with pytest.raises(ValueError, match=".*one event.*"):
-        _compute_time_to_event_metric(
-            metric_func=metric_func,
+        _prep_data_for_time_to_event_metric(
             train_data=Mock(predictive=Mock(targets=Mock(num_features=999))),
             test_data=mock_test_data,
             horizons=Mock(),
@@ -330,8 +320,7 @@ def test_compute_time_to_event_metric():
         )
 
     with pytest.raises(ValueError, match=".*targets.*"):
-        _compute_time_to_event_metric(
-            metric_func=metric_func,
+        _prep_data_for_time_to_event_metric(
             train_data=Mock(predictive=Mock(targets=None)),
             test_data=mock_test_data,
             horizons=Mock(),
@@ -339,8 +328,7 @@ def test_compute_time_to_event_metric():
         )
 
     with pytest.raises(ValueError, match=".*equal number of time steps.*"):
-        _compute_time_to_event_metric(
-            metric_func=metric_func,
+        _prep_data_for_time_to_event_metric(
             train_data=mock_train_data,
             test_data=mock_test_data,
             horizons=[1, 2, 3],
diff --git a/tests/core/test_core_utils.py b/tests/core/test_core_utils.py
index 59a1f0ce..a93cf077 100644
--- a/tests/core/test_core_utils.py
+++ b/tests/core/test_core_utils.py
@@ -1,3 +1,5 @@
+# pylint: disable=unnecessary-pass
+
 import re
 from typing import Callable
 
@@ -46,6 +48,30 @@ def test_fails_difference(self, lit, d, expect):
             utils.ensure_literal_matches_dict_keys(lit, d)
 
 
+# Test RichReprStrPassthrough ----------
+
+
+def test_rich_repr_str_passthrough():
+    rich_repr_passthrough = utils.RichReprStrPassthrough("foobar")
+    assert rich_repr_passthrough.string == "foobar"
+    assert repr(rich_repr_passthrough) == "foobar"
+
+
+# Test RichReprStrPassthrough (end) ----------
+
+
+# Test is_iterable ----------
+
+
+def test_is_iterable():
+    assert utils.is_iterable([1, 2, 3])
+    assert utils.is_iterable("abc")
+    assert utils.is_iterable(999) is False
+
+
+# Test is_iterable (end) ----------
+
+
 # Test get_from_args_or_kwargs ----------
 
 
@@ -261,6 +287,9 @@ def test_builtin(self):
 # Test get_class_full_name (end) ----------
 
 
+# Test unique_in_order_of_appearance ----------
+
+
 class TestUniqueInOrderOfAppearance:
     def test_strs(self):
         assert utils.unique_in_order_of_appearance(["a", "a", "b", "c", "c", "a", "b"]) == ["a", "b", "c"]
@@ -270,3 +299,145 @@ def test_ints(self):
 
     def test_mixed(self):
         assert utils.unique_in_order_of_appearance([1, "a", 2, "a", 2, 3, 3, 1, 1, 2]) == [1, "a", 2, 3]
+
+
+# Test unique_in_order_of_appearance (end) ----------
+
+
+# Test is_method_defined_in_class ----------
+
+
+class Parent:
+    def __init__(self):
+        pass
+
+    def some_method(self):
+        pass
+
+
+class ChildInherits(Parent):
+    pass
+
+
+class ChildDoesNotInherit(Parent):
+    def __init__(self) -> None:
+        pass
+
+    def some_method(self):
+        pass
+
+
+class TestIsMethodDefinedInClass:
+    def test_yes_class(self):
+        assert utils.is_method_defined_in_class(ChildDoesNotInherit, "__init__")
+        assert utils.is_method_defined_in_class(ChildDoesNotInherit, "some_method")
+
+    def test_no_class(self):
+        assert utils.is_method_defined_in_class(ChildInherits, "__init__") is False
+        assert utils.is_method_defined_in_class(ChildInherits, "some_method") is False
+
+    def test_yes_instance(self):
+        assert utils.is_method_defined_in_class(ChildDoesNotInherit(), "__init__")
+        assert utils.is_method_defined_in_class(ChildDoesNotInherit(), "some_method")
+
+    def test_no_instance(self):
+        assert utils.is_method_defined_in_class(ChildInherits(), "__init__") is False
+        assert utils.is_method_defined_in_class(ChildInherits(), "some_method") is False
+
+
+# Test is_method_defined_in_class (end) ----------
+
+
+# Test clean_multiline_docstr ----------
+
+
+class ClassWithMultilineDocstr:
+    """This is a docstring.
+    It has multiple lines.
+    """
+
+    pass
+
+
+class ClassWithSingleLineDocstr:
+    """This is a docstring."""
+
+    pass
+
+
+class TestCleanMultilineDocstr:
+    def test_multiline(self):
+        assert (
+            utils.clean_multiline_docstr(ClassWithMultilineDocstr.__doc__)  # pyright: ignore
+            == "This is a docstring. It has multiple lines."
+        )
+
+    def test_single_line(self):
+        assert (
+            utils.clean_multiline_docstr(ClassWithSingleLineDocstr.__doc__) == "This is a docstring."  # pyright: ignore
+        )
+
+
+# Test clean_multiline_docstr (end) ----------
+
+
+# Test make_description_from_doc ----------
+
+
+class NoDocstr:
+    pass
+
+
+class DocstrClassOnly:
+    """This is a docstring."""
+
+    pass
+
+
+class DocstrInitOnly:
+    def __init__(self) -> None:
+        """This is an __init__ docstring."""
+        pass
+
+
+class DocstrClassAndInit:
+    """This is a docstring."""
+
+    def __init__(self) -> None:
+        """This is an __init__ docstring."""
+        pass
+
+
+class ParentWithInit:
+    def __init__(self) -> None:
+        """This is an __init__ docstring but it's on parent."""
+        pass
+
+
+class DocstrClassAndInitButInitOnlyFromParent:
+    """This is a docstring."""
+
+    pass
+
+
+class TestMakeDescriptionFromDoc:
+    def test_docstr_class_only(self):
+        assert utils.make_description_from_doc(DocstrClassOnly()) == "This is a docstring."
+
+    def test_init_only(self):
+        assert utils.make_description_from_doc(DocstrInitOnly()) == "This is an __init__ docstring."
+
+    def test_docstr_class_and_init(self):
+        assert (
+            utils.make_description_from_doc(DocstrClassAndInit())
+            == "This is a docstring. This is an __init__ docstring."
+        )
+
+    def test_docstr_class_and_init_but_init_only_from_parent(self):
+        assert utils.make_description_from_doc(DocstrClassAndInitButInitOnlyFromParent()) == "This is a docstring."
+
+    def test_docstr_truncated(self):
+        assert utils.make_description_from_doc(DocstrClassOnly(), max_len_keep=5) == "This ..."
+
+
+# Test make_description_from_doc (end) ----------
diff --git a/tests/metrics/prediction/one_off/test_metrics_prediction_one_off_builtin_classification.py b/tests/metrics/prediction/one_off/test_metrics_prediction_one_off_builtin_classification.py
new file mode 100644
index 00000000..0b285280
--- /dev/null
+++ b/tests/metrics/prediction/one_off/test_metrics_prediction_one_off_builtin_classification.py
@@ -0,0 +1,124 @@
+# pylint: disable=protected-access
+
+import numpy as np
+import pytest
+
+from tempor import plugin_loader
+from tempor.metrics.prediction.one_off import plugin_builtin_classification
+
+# ------------------------------------------------------------------------------
+# Test utilities.
+
+
+class TestUtilities:
+    class TestCastToYPred:
+        def test_with_probabilities(self):
+            y_pred_proba = np.array([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]])
+            expected = np.array([1, 0, 1])
+            result = plugin_builtin_classification._cast_to_y_pred(y_pred_proba)
+            np.testing.assert_array_equal(result, expected)
+
+        def test_with_class_labels(self):
+            y_pred_proba = np.array([1, 0, 1])
+            expected = np.array([1, 0, 1])
+            result = plugin_builtin_classification._cast_to_y_pred(y_pred_proba)
+            np.testing.assert_array_equal(result, expected)
+
+        def test_with_invalid_dimensions(self):
+            y_pred_proba = np.array([[[0.1, 0.9], [0.8, 0.2]]])
+            with pytest.raises(ValueError):
+                plugin_builtin_classification._cast_to_y_pred(y_pred_proba)
+
+    def test_prep_auc_multiclass_fails(self):
+        with pytest.raises(ValueError, match=".*NaN.*"):
+            plugin_builtin_classification._prep_auc_multiclass(np.array([1, 0]), np.array([[0.1, 0.9], [np.nan, 0.2]]))
+
+    class TestGetYPredProbaHlpr:
+        def test_2class_2prob(self):
+            in_ = np.array([[0.1, 0.9], [0.3, 0.7], [0.11, 0.89]])
+            exp = np.array([[0.9], [0.7], [0.89]]).ravel()
+            out = plugin_builtin_classification._get_y_pred_proba_hlpr(in_, nclasses=2)
+            assert (out == exp).all()
+
+        def test_2class_1prob(self):
+            in_ = np.array([[0.1], [0.3], [0.11]])
+            out = plugin_builtin_classification._get_y_pred_proba_hlpr(in_, nclasses=2)
+            assert (out == in_).all()
+
+
+# ------------------------------------------------------------------------------
+# Test the actual metrics.
+
+
+TEST_CASE_PAIRS = [
+    # Cases: class labels.
+    (  # 1. Perfect prediction:
+        np.array([1, 0, 1, 0]),
+        np.array([1, 0, 1, 0]),
+    ),
+    (  # 2. Completely incorrect prediction:
+        np.array([1, 0, 1, 0]),
+        np.array([0, 1, 0, 1]),
+    ),
+    (  # 3. Partially correct prediction:
+        np.array([1, 0, 1, 0]),
+        np.array([1, 1, 0, 0]),
+    ),
+    # Cases: class probabilities.
+    (  # 4. Perfect prediction:
+        np.array([1, 0, 1, 0]),
+        np.array([[0.1, 0.9], [0.7, 0.3], [0.2, 0.8], [0.6, 0.4]]),
+    ),
+    (  # 5. Completely incorrect prediction:
+        np.array([1, 0, 1, 0]),
+        np.array([[0.9, 0.1], [0.3, 0.7], [0.8, 0.2], [0.4, 0.6]]),
+    ),
+    (  # 6. Partially correct prediction:
+        np.array([1, 0, 1, 0]),
+        np.array([[0.1, 0.9], [0.3, 0.7], [0.8, 0.2], [0.6, 0.4]]),
+    ),
+    # Cases: 2+ classes.
+    (  # 7. Perfect prediction:
+        np.array([1, 0, 2, 0]),
+        np.array([[0.1, 0.9, 0.0], [0.7, 0.3, 0.0], [0.0, 0.0, 1.0], [0.6, 0.4, 0.0]]),
+    ),
+    (  # 8. 1 of 4 is correct prediction:
+        np.array([1, 0, 2, 0]),
+        np.array([[0.9, 0.1, 0.0], [0.3, 0.7, 0.0], [0.0, 0.0, 1.0], [0.4, 0.6, 0.0]]),
+    ),
+    # NOTE: Any additional cases here.
+]
+
+EXPECTED_VALUES = {
+    "accuracy": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.25],
+    "f1_score_micro": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.25],
+    "f1_score_macro": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.3333],
+    "f1_score_weighted": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.25],
+    "kappa": [1.0, -1.0, 0.0, 1.0, -1.0, 0.0, 1.0, -0.0909],
+    "kappa_quadratic": [1.0, -1.0, 0.0, 1.0, -1.0, 0.0, 1.0, 0.4],
+    "recall_micro": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.25],
+    "recall_macro": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.3333],
+    "recall_weighted": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.25],
+    "precision_micro": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.25],
+    "precision_macro": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.3333],
+    "precision_weighted": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.25],
+    "mcc": [1.0, -1.0, 0.0, 1.0, -1.0, 0.0, 1.0, -0.1],
+    "aucprc": [1.0, 0.5, 0.5, 1.0, 0.41666, 0.75, 1.0, 0.61785],
+    "aucroc": [1.0, 0.0, 0.5, 1.0, 0.0, 0.5, 1.0, 0.71875],
+    # NOTE: Any additional metrics here.
+    # "<metric_name>": [...expected values corresponding to the test cases above...]
+}
+
+METRIC_NAMES = list(EXPECTED_VALUES.keys())
+
+
+@pytest.mark.parametrize("metric_name", METRIC_NAMES)
+@pytest.mark.parametrize("case_idx, actual, predicted", [(idx, a, p) for idx, (a, p) in enumerate(TEST_CASE_PAIRS)])
+def test_accuracy_metric(metric_name, case_idx, actual, predicted):
+    metric = plugin_loader.get(f"prediction.one_off.classification.{metric_name}", plugin_type="metric")
+    result = metric.evaluate(actual, predicted)
+
+    expected = EXPECTED_VALUES[metric_name][case_idx]
+
+    assert metric.direction in ("maximize", "minimize")
+    np.testing.assert_almost_equal(result, expected, 4)
diff --git a/tests/metrics/prediction/one_off/test_metrics_prediction_one_off_builtin_regression.py b/tests/metrics/prediction/one_off/test_metrics_prediction_one_off_builtin_regression.py
new file mode 100644
index 00000000..dbe87e43
--- /dev/null
+++ b/tests/metrics/prediction/one_off/test_metrics_prediction_one_off_builtin_regression.py
@@ -0,0 +1,38 @@
+import numpy as np
+import pytest
+
+from tempor import plugin_loader
+
+TEST_CASE_PAIRS = [
+    (  # 1. Perfect prediction:
+        np.array([2.0, 7.5, 2.1, 9.3]),
+        np.array([2.0, 7.5, 2.1, 9.3]),
+    ),
+    (  # 3. Non-perfect prediction:
+        np.array([2.0, 7.5, 2.1, 9.3]),
+        np.array([-1.2, 7.1, 2.8, 11.2]),
+    ),
+    # NOTE: Any additional cases here.
+]
+
+EXPECTED_VALUES = {
+    "mse": [0.0, 3.6249],
+    "mae": [0.0, 1.5499],
+    "r2": [1.0, 0.6543],
+    # NOTE: Any additional metrics here.
+    # "<metric_name>": [...expected values corresponding to the test cases above...]
+}
+
+METRIC_NAMES = list(EXPECTED_VALUES.keys())
+
+
+@pytest.mark.parametrize("metric_name", METRIC_NAMES)
+@pytest.mark.parametrize("case_idx, actual, predicted", [(idx, a, p) for idx, (a, p) in enumerate(TEST_CASE_PAIRS)])
+def test_accuracy_metric(metric_name, case_idx, actual, predicted):
+    metric = plugin_loader.get(f"prediction.one_off.regression.{metric_name}", plugin_type="metric")
+    result = metric.evaluate(actual, predicted)
+
+    expected = EXPECTED_VALUES[metric_name][case_idx]
+
+    assert metric.direction in ("maximize", "minimize")
+    np.testing.assert_almost_equal(result, expected, 4)
diff --git a/tests/metrics/test_metrics.py b/tests/metrics/test_metrics.py
new file mode 100644
index 00000000..c8b31539
--- /dev/null
+++ b/tests/metrics/test_metrics.py
@@ -0,0 +1,54 @@
+import re
+from typing import Any
+
+import pytest
+
+from tempor.metrics.metric import Metric
+from tempor.metrics.metric_typing import MetricDirection
+
+
+class MyMetric(Metric):
+    """My metric under test."""
+
+    name = "my_metric"
+    category = "my_category"
+    plugin_type = "metric"
+
+    @property
+    def direction(self) -> MetricDirection:
+        return "maximize"
+
+    def _evaluate(self, actual: Any, predicted: Any, *args: Any, **kwargs: Any) -> Any:
+        return 0.5
+
+
+class TestMetric:
+    def test_init_success(self):
+        metric = MyMetric()
+        assert metric.name == "my_metric"
+        assert metric.direction == "maximize"
+
+    def test_evaluate_success(self):
+        metric = MyMetric()
+        actual = [1.0, 2.0, 3.0]
+        predicted = [1.0, 2.0, 3.0]
+        score = metric.evaluate(actual=actual, predicted=predicted)
+        assert score == 0.5
+
+    def test_evaluate_via_call_success(self):
+        metric = MyMetric()
+        actual = [1.0, 2.0, 3.0]
+        predicted = [1.0, 2.0, 3.0]
+        score = metric(actual=actual, predicted=predicted)
+        assert score == 0.5
+
+    def test_validation_fails(self):
+        metric = MyMetric()
+        with pytest.raises(ValueError):
+            metric.evaluate(actual=None, predicted=[1.0, 2.0, 3.0])
+        with pytest.raises(ValueError):
+            metric.evaluate(actual=[1.0, 2.0, 3.0], predicted=None)
+
+    def test_repr(self):
+        metric = MyMetric()
+        assert re.match(r"MyMetric\(name='my_metric', description='My metric under test.'\)", repr(metric)) is not None
diff --git a/tests/metrics/time_to_event/test_metrics_time_to_event_builtin.py b/tests/metrics/time_to_event/test_metrics_time_to_event_builtin.py
new file mode 100644
index 00000000..825e66a8
--- /dev/null
+++ b/tests/metrics/time_to_event/test_metrics_time_to_event_builtin.py
@@ -0,0 +1,69 @@
+from typing import NamedTuple, Optional
+
+import numpy as np
+import pytest
+
+from tempor import plugin_loader
+from tempor.data import data_typing
+from tempor.metrics import metric_typing
+
+
+class TestData(NamedTuple):
+    actual: metric_typing.EventArrayTimeArray
+    predicted: np.ndarray
+    horizons: data_typing.TimeIndex
+    actual_train: Optional[metric_typing.EventArrayTimeArray] = None
+
+
+TEST_CASE_PAIRS = [
+    TestData(
+        actual=(
+            np.asarray([1, 0, 1, 1, 1, 0]),
+            np.asarray([100.0, 110.0, 108.0, 99.0, 101.0, 140.0]),
+        ),
+        # predicted: (n_samples, n_horizons_timesteps, n_features)
+        predicted=np.expand_dims(
+            np.asarray(
+                [
+                    [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
+                    [0.2, 0.3, 0.4, 0.5, 0.6, 0.7],
+                    [0.3, 0.4, 0.5, 0.6, 0.7, 0.8],
+                ]
+            ).T,
+            axis=-1,
+        ),
+        horizons=[100.0, 115.0, 130.0],
+        actual_train=(
+            np.asarray([1, 0, 1, 0, 0, 0, 1, 1, 0, 1]),
+            np.asarray([100.0, 102.0, 107.0, 120.0, 103.0, 133.0, 114.0, 170.0, 113.0, 124.0]),
+        ),
+    ),
+    # NOTE: Any additional cases here.
+]
+
+EXPECTED_VALUES = {
+    "c_index": [
+        [0.6, 0.4346, 0.4346],
+    ],
+    "brier_score": [
+        [0.285, 0.1657, 0.2239],
+    ],
+    # NOTE: Any additional metrics here.
+    # "<metric_name>": [...expected arrays (metric value per horizon) corresponding to the test cases above...]
+}
+
+METRIC_NAMES = list(EXPECTED_VALUES.keys())
+
+
+@pytest.mark.parametrize("metric_name", METRIC_NAMES)
+@pytest.mark.parametrize("case_idx, case_data", enumerate(TEST_CASE_PAIRS))
+def test_accuracy_metric(metric_name, case_idx, case_data):
+    metric = plugin_loader.get(f"time_to_event.{metric_name}", plugin_type="metric")
+
+    actual, predicted, horizons, actual_train = case_data
+    result = metric.evaluate(actual, predicted, horizons, actual_train=actual_train)
+
+    expected = EXPECTED_VALUES[metric_name][case_idx]
+
+    assert metric.direction in ("maximize", "minimize")
+    np.testing.assert_almost_equal(result, expected, 4)
diff --git a/tests/benchmarks/test_metrics.py b/tests/metrics/time_to_event/test_metrics_time_to_event_metric_impl.py
similarity index 73%
rename from tests/benchmarks/test_metrics.py
rename to tests/metrics/time_to_event/test_metrics_time_to_event_metric_impl.py
index 2f74d3eb..7e054826 100644
--- a/tests/benchmarks/test_metrics.py
+++ b/tests/metrics/time_to_event/test_metrics_time_to_event_metric_impl.py
@@ -5,20 +5,20 @@
 import numpy as np
 import pytest
 
-from tempor.benchmarks import metrics
+from tempor.metrics.time_to_event import _metric_impl
 
 
 def test_validation_check_y_survival():
     with pytest.raises(ValueError, match=".*y must be.*"):
-        metrics.check_y_survival(y_or_event="wrong")  # type: ignore
+        _metric_impl.check_y_survival(y_or_event="wrong")  # type: ignore
 
     with pytest.raises(ValueError, match=".*boolean.*"):
-        metrics.check_y_survival(
+        _metric_impl.check_y_survival(
             y_or_event=np.asarray([[1, 3, 7], [1, 1, 1]], dtype={"names": ("a", "b"), "formats": ("int", "f8")})
         )
 
     with pytest.raises(ValueError, match=".*all.*censored.*"):
-        metrics.check_y_survival(
+        _metric_impl.check_y_survival(
             y_or_event=np.asarray(
                 [[False, False, False], [0, 0, 0]], dtype={"names": ("a", "b"), "formats": ("bool", "f8")}
             ),
@@ -28,38 +28,40 @@ def test_validation_check_y_survival():
 
 def test_validation_check_estimate_1d():
     with pytest.raises(ValueError, match=".*Expected 1D.*"):
-        metrics._check_estimate_1d([[1, 1], [2, 2]], test_time=100)
+        _metric_impl._check_estimate_1d([[1, 1], [2, 2]], test_time=100)  # type: ignore
 
 
 def test_validation_check_times():
     with pytest.raises(ValueError, match=".*within.*"):
-        metrics._check_times(np.asarray([1000]), [1, 4])
+        _metric_impl._check_times(np.asarray([1000]), [1, 4])  # type: ignore
 
 
 def test_validation_check_estimate_2d():
     with pytest.raises(ValueError, match=".*estimate with.*"):
-        metrics._check_estimate_2d(np.asarray([[1, 2], [1, 2]]), np.asarray([1, 3]), np.asarray([1]), Mock())
+        _metric_impl._check_estimate_2d(np.asarray([[1, 2], [1, 2]]), np.asarray([1, 3]), np.asarray([1]), Mock())
 
 
 def test_compute_counts():
-    out = metrics._compute_counts(np.asarray([True, True, False]), np.asarray([1, 8, 9]))
+    out = _metric_impl._compute_counts(np.asarray([True, True, False]), np.asarray([1, 8, 9]))
     assert len(out) == 4
 
 
 def test_compute_counts_order_given():
-    out = metrics._compute_counts(np.asarray([True, True, False]), np.asarray([1, 8, 9]), order=np.asarray([2, 1, 0]))
+    out = _metric_impl._compute_counts(
+        np.asarray([True, True, False]), np.asarray([1, 8, 9]), order=np.asarray([2, 1, 0])
+    )
     assert len(out) == 4
 
 
 def test_compute_counts_truncated_validation():
     with pytest.raises(ValueError, match=".*must be larger.*"):
-        metrics._compute_counts_truncated(
+        _metric_impl._compute_counts_truncated(
             np.asarray([True, True, False]), np.asarray([11, 11, 11]), np.asarray([1, 1, 1])
         )
 
 
 def test_kaplan_meier_estimator_reverse():
-    uniq_times, y = metrics.kaplan_meier_estimator(
+    uniq_times, y = _metric_impl.kaplan_meier_estimator(
         event=np.asarray([True, True, False]),
         time_exit=np.asarray([1, 8, 9]),
         reverse=True,
@@ -69,7 +71,7 @@ def test_kaplan_meier_estimator_reverse():
 
 
 def test_kaplan_meier_estimator_time_enter():
-    uniq_times, y = metrics.kaplan_meier_estimator(
+    uniq_times, y = _metric_impl.kaplan_meier_estimator(
         event=np.asarray([True, True, False]),
         time_exit=np.asarray([1, 8, 9]),
         time_enter=np.asarray([0, 1, 0]),
@@ -78,7 +80,7 @@ def test_kaplan_meier_estimator_time_enter():
     assert len(y) == 4
 
     with pytest.raises(ValueError, match=".*left.*"):
-        metrics.kaplan_meier_estimator(
+        _metric_impl.kaplan_meier_estimator(
             event=np.asarray([True, True, False]),
             time_exit=np.asarray([1, 8, 9]),
             time_enter=np.asarray([0, 1, 0]),
@@ -87,7 +89,7 @@ def test_kaplan_meier_estimator_time_enter():
 
 
 def test_kaplan_meier_estimator_time_min():
-    uniq_times, y = metrics.kaplan_meier_estimator(
+    uniq_times, y = _metric_impl.kaplan_meier_estimator(
         event=np.asarray([True, True, False]),
         time_exit=np.asarray([1, 8, 9]),
         time_min=0,
@@ -97,7 +99,7 @@ def test_kaplan_meier_estimator_time_min():
 
 
 def test_survival_function_estimator_fit():
-    model = metrics.SurvivalFunctionEstimator()
+    model = _metric_impl.SurvivalFunctionEstimator()
     model.fit(
         y=np.asarray(
             list(zip([False, False, True, True], [1, 2, 11, 4])), dtype={"names": ("e", "t"), "formats": ("bool", "f8")}
@@ -106,7 +108,7 @@ def test_survival_function_estimator_fit():
 
 
 def test_survival_function_estimator_predict_proba_invalid():
-    model = metrics.SurvivalFunctionEstimator()
+    model = _metric_impl.SurvivalFunctionEstimator()
     model.fit(
         y=np.asarray(
             list(zip([False, False, True, True], [1, 2, 11, 4])), dtype={"names": ("e", "t"), "formats": ("bool", "f8")}
@@ -118,7 +120,7 @@ def test_survival_function_estimator_predict_proba_invalid():
 
 
 def test_censoring_distribution_estimator_fit_all():
-    model = metrics.CensoringDistributionEstimator()
+    model = _metric_impl.CensoringDistributionEstimator()
     model.fit(
         y=np.asarray(
             list(zip([True, True, True, True], [1, 2, 11, 4])), dtype={"names": ("e", "t"), "formats": ("bool", "f8")}
@@ -127,7 +129,7 @@ def test_censoring_distribution_estimator_fit_all():
 
 
 def test_censoring_distribution_estimator_predict_ipcw_ghat_0(monkeypatch):
-    model = metrics.CensoringDistributionEstimator()
+    model = _metric_impl.CensoringDistributionEstimator()
     monkeypatch.setattr(model, "predict_proba", Mock(return_value=np.asarray([0, 0, 0, 0])))
 
     model.fit(
@@ -145,31 +147,31 @@ def test_censoring_distribution_estimator_predict_ipcw_ghat_0(monkeypatch):
 
 
 def test_estimate_concordance_index_invalid_comparable(monkeypatch):
-    monkeypatch.setattr(metrics, "_get_comparable", Mock(return_value=([], [])))
+    monkeypatch.setattr(_metric_impl, "_get_comparable", Mock(return_value=([], [])))
 
     with pytest.raises(ValueError, match=".*pairs.*"):
-        metrics._estimate_concordance_index(Mock(), Mock(), Mock(), Mock())
+        _metric_impl._estimate_concordance_index(Mock(), Mock(), Mock(), Mock())
 
 
 def test_create_structured_array_validation():
     with pytest.raises(ValueError, match=".*different.*"):
-        metrics.create_structured_array([], [], name_event="a", name_time="a")
+        _metric_impl.create_structured_array([], [], name_event="a", name_time="a")
 
 
 def test_create_structured_array_non_bool_event():
-    y = metrics.create_structured_array([1, 1, 0, 1], [22, 34, 22, 134])
+    y = _metric_impl.create_structured_array([1, 1, 0, 1], [22, 34, 22, 134])
     assert len(y) == 4
 
 
 def test_create_structured_array_non_bool_event_validation():
     with pytest.raises(ValueError, match=".*binary.*"):
-        metrics.create_structured_array([1, 22, 1111, 44544], [22, 34, 22, 134])
+        _metric_impl.create_structured_array([1, 22, 1111, 44544], [22, 34, 22, 134])
     with pytest.raises(ValueError, match=".*0 and 1.*"):
-        metrics.create_structured_array([3, 4, 3, 4], [22, 34, 22, 134])
+        _metric_impl.create_structured_array([3, 4, 3, 4], [22, 34, 22, 134])
 
 
 def test_brier_score_reshape():
-    out = metrics.brier_score(
+    out = _metric_impl.brier_score(
         survival_train=np.asarray(
             list(zip([True, True, True, True], [1, 2, 11, 4])),
             dtype={"names": ("e", "t"), "formats": ("bool", "f8")},
@@ -185,7 +187,7 @@ def test_brier_score_reshape():
 
 
 def test_concordance_index_ipcw_tau_none():
-    out = metrics.concordance_index_ipcw(
+    out = _metric_impl.concordance_index_ipcw(
         survival_train=np.asarray(
             list(zip([True, True, True, True], [1, 2, 11, 4])),
             dtype={"names": ("e", "t"), "formats": ("bool", "f8")},
diff --git a/tests/nb_eval.py b/tests/nb_eval.py
index 124e7a46..61f9ab23 100644
--- a/tests/nb_eval.py
+++ b/tests/nb_eval.py
@@ -56,19 +56,20 @@ def filter_notebooks(all_paths: List, enabled_tests: List, excluded_dirs: List)
     return notebooks
 
 
-def run_notebook(notebook_path: Path, skip_pip_install: bool = True) -> None:
+def run_notebook(notebook_path: Path, skip_pip_install: bool = True, overwrite_file: bool = False) -> None:
     """Run a notebook.
 
     Args:
         notebook_path (Path): Path to notebook.
-        skip_pip_install (bool, optional): Comment out the pip install commands in cells if present. Defaults to True.
+        skip_pip_install (bool, optional): Comment out the pip install commands in cells if present. Defaults to `True`.
+        overwrite_file (bool, optional): Whether to overwrite the notebook file. Defaults to `False`.
     """
     with open(notebook_path, encoding="utf8") as f:
         f_str = f.read()
 
     if skip_pip_install:
         # Comment out `pip install` commands, as do not need to run those in test.
-        f_str = re.sub(r"[%!]\s*pip", "#pip", f_str, flags=re.DOTALL)
+        f_str = re.sub(r"[%!]\s*pip", "#%pip", f_str, flags=re.DOTALL)
 
     nb = nbformat.reads(f_str, as_version=4)
     proc = ExecutePreprocessor(timeout=6000)
@@ -76,13 +77,38 @@ def run_notebook(notebook_path: Path, skip_pip_install: bool = True) -> None:
     # Will raise on cell error:
     proc.preprocess(nb, {"metadata": {"path": workspace}})
 
+    if overwrite_file:
+        # Cell modifications to avoid "volatile" metadata causing unnecessary diffs.
+        for idx, cell in enumerate(nb.cells):
+            # Clear the metadata section.
+            if hasattr(cell, "metadata"):
+                cell.metadata = {}
+            # Clear `execution_count` if present.
+            if cell.cell_type == "code" and hasattr(cell, "execution_count"):
+                cell.execution_count = None
+            # Clear cell outputs' `execution_count` if present.
+            if hasattr(cell, "outputs"):
+                for output in cell.outputs:
+                    if hasattr(output, "execution_count"):
+                        output.execution_count = None
+            # Restore pip install commands if present.
+            if skip_pip_install and cell.cell_type == "code":
+                if "#%pip" in cell.source:
+                    cell.source = cell.source.replace("#%pip", "%pip")
+            # Overwrite the cell:
+            nb["cells"][idx] = cell
+
+        with open(notebook_path, "w", encoding="utf8") as f:
+            nbformat.write(nb, f)
+
 
 @joblib.delayed
-def test_notebook(p: Path) -> None:
+def test_notebook(p: Path, overwrite_file: bool = False) -> None:
     """Test a notebook.
 
     Args:
         p (Path): Path to notebook.
+        overwrite_file (bool, optional): Whether to overwrite the notebook file. Defaults to `False`.
 
     Raises:
         BaseException: Exception raised if notebook fails.
@@ -90,7 +116,7 @@ def test_notebook(p: Path) -> None:
     print("Testing notebook:", p.name)
     start = time.time()
     try:
-        run_notebook(p)
+        run_notebook(p, skip_pip_install=True, overwrite_file=overwrite_file)
     except BaseException as e:
         print("FAIL", p.name, e)
         raise e
@@ -99,18 +125,50 @@ def test_notebook(p: Path) -> None:
 
 
 @click.command()
-@click.option("--nb_dir", type=str, default=".", help="Path to directory containing notebooks to test")
-@click.option("-n", "--n_jobs", type=int, default=-1, help="Number of joblib Parallel jobs")
-@click.option("-v", "--verbose", type=int, default=0, help="Verbosity value for joblib Parallel")
-def main(nb_dir: Path, n_jobs: int, verbose: int) -> None:
+@click.option(
+    "--nb_dir",
+    type=str,
+    default=".",
+    help="Path to directory containing notebooks to test.",
+    show_default=True,
+)
+@click.option(
+    "-n",
+    "--n_jobs",
+    type=int,
+    default=-1,
+    help="Number of joblib Parallel jobs.",
+    show_default=True,
+)
+@click.option(
+    "-v",
+    "--verbose",
+    type=int,
+    default=0,
+    help="Verbosity value for joblib Parallel.",
+    show_default=True,
+)
+@click.option(
+    "-o",
+    "--overwrite_notebook",
+    is_flag=True,
+    default=False,
+    help="Whether to overwrite the notebooks with executed versions.",
+    show_default=True,
+)
+def main(nb_dir: Path, n_jobs: int, verbose: int, overwrite_notebook: bool) -> None:
     """Test all notebooks in a directory.
 
     Args:
         nb_dir (Path): Path to directory containing notebooks to test.
         n_jobs (int): Number of joblib Parallel jobs.
         verbose (int): Verbosity value for joblib Parallel.
+        overwrite_notebook (bool): Whether to overwrite the notebooks with executed versions.
     """
 
+    if overwrite_notebook:
+        print("WARNING: Overwriting notebooks with executed versions as --overwrite_notebook was enabled.")
+
     nb_dir = Path(nb_dir)
     notebook_paths = filter_notebooks(list(nb_dir.rglob("*")), ENABLED_TESTS, EXCLUDED_DIRS)
     print("Notebooks to be tested:")
@@ -118,7 +176,9 @@ def main(nb_dir: Path, n_jobs: int, verbose: int) -> None:
 
     print(f"Running n_jobs={n_jobs} parallel jobs. Note: -1 means as many jobs as available CPUs.")
     start = time.time()
-    joblib.Parallel(n_jobs=n_jobs, backend="loky", verbose=verbose)(test_notebook(p) for p in notebook_paths)
+    joblib.Parallel(n_jobs=n_jobs, backend="loky", verbose=verbose)(
+        test_notebook(p, overwrite_file=overwrite_notebook) for p in notebook_paths
+    )
     print(f"Testing all notebooks took {time.time() - start:.2f} sec")
 
 
diff --git a/tests/test_plugin_loader.py b/tests/test_plugin_loader.py
index a39e63ce..9619b2b9 100644
--- a/tests/test_plugin_loader.py
+++ b/tests/test_plugin_loader.py
@@ -104,11 +104,61 @@ def test_datasource_contents():
     assert "dummy_treatments" in all_plugins["treatments"]["temporal"]
 
 
+def test_metrics_contents():
+    all_plugins = plugin_loader.list(plugin_type="metric")
+
+    # Check categories:
+    assert "prediction" in all_plugins
+    assert "time_to_event" in all_plugins
+    # assert "treatments" in all_plugins
+
+    # Check subcategories:
+    assert "one_off" in all_plugins["prediction"]
+    # assert "temporal" in all_plugins["prediction"]
+    # assert "one_off" in all_plugins["treatments"]
+    # assert "temporal" in all_plugins["treatments"]
+
+    # Check sub-subcategories:
+    assert "classification" in all_plugins["prediction"]["one_off"]
+    assert "regression" in all_plugins["prediction"]["one_off"]
+    # assert "classification" in all_plugins["prediction"]["temporal"]
+    # assert "regression" in all_plugins["prediction"]["temporal"]
+    # assert "classification" in all_plugins["treatments"]["one_off"]
+    # assert "regression" in all_plugins["treatments"]["one_off"]
+    # assert "classification" in all_plugins["treatments"]["temporal"]
+    # assert "regression" in all_plugins["treatments"]["temporal"]
+
+    # Check plugins:
+    assert "accuracy" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "f1_score_micro" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "f1_score_macro" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "f1_score_weighted" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "kappa" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "kappa_quadratic" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "recall_micro" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "recall_macro" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "recall_weighted" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "precision_micro" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "precision_macro" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "precision_weighted" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "mcc" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "aucprc" in all_plugins["prediction"]["one_off"]["classification"]
+    assert "aucroc" in all_plugins["prediction"]["one_off"]["classification"]
+    # ---
+    assert "mse" in all_plugins["prediction"]["one_off"]["regression"]
+    assert "mae" in all_plugins["prediction"]["one_off"]["regression"]
+    assert "r2" in all_plugins["prediction"]["one_off"]["regression"]
+    # ---
+    assert "c_index" in all_plugins["time_to_event"]
+    assert "brier_score" in all_plugins["time_to_event"]
+
+
 def test_plugin_types():
     plugin_types = plugin_loader.list_plugin_types()
 
     assert "method" in plugin_types
     assert "datasource" in plugin_types
+    assert "metric" in plugin_types
 
 
 PLUGIN_FQNS = plugin_loader.list_full_names()
diff --git a/tutorials/extending/tutorial02_testing_custom_method.ipynb b/tutorials/extending/tutorial02_testing_custom_method.ipynb
index 5db41f44..4faa2453 100644
--- a/tutorials/extending/tutorial02_testing_custom_method.ipynb
+++ b/tutorials/extending/tutorial02_testing_custom_method.ipynb
@@ -65,7 +65,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -417,8 +417,7 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.18"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2