diff --git a/docs/source/tutorials/tslib_v2_example.ipynb b/docs/source/tutorials/tslib_v2_example.ipynb
index 1a4d8ad94..0f3c8394a 100644
--- a/docs/source/tutorials/tslib_v2_example.ipynb
+++ b/docs/source/tutorials/tslib_v2_example.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "b5d44943",
+   "id": "66a17c35",
    "metadata": {},
    "source": [
     "# TSLib for v2 - Example notebook for full pipeline"
@@ -10,51 +10,47 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b7d27b55",
+   "id": "2e986d11",
    "metadata": {},
    "source": [
-    "## Basic imports for getting started\n",
+    "<div class=\"alert alert-block alert-warning\">\n",
     "\n",
-    "This notebook is a basic vignette for the usage of the `tslib` data module on the `TimeXer` model for the v2 of PyTorch Forecasting. This is an experimental version and is an unstable version of the API.\n",
+    ":warning: <b>Experimental, unstable API.</b> The <code>tslib</code> models and their data layer are part of an experimental rework of <code>pytorch-forecasting</code>, planned for the v2.0.0 release. The API is unstable and may change without prior notice, and is <b>not recommended for production use</b>. Feedback is very welcome on <a href=\"https://github.com/sktime/pytorch-forecasting/pull/1836\">issue&nbsp;#1836</a>.\n",
     "\n",
-    "Feedback and suggestions on this pipeline - PR [#1836](https://github.com/sktime/pytorch-forecasting/pull/1836)"
+    "</div>"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "550a3fbf",
+   "cell_type": "markdown",
+   "id": "405d0acc",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "import pandas as pd\n",
-    "from sklearn.preprocessing import StandardScaler\n",
-    "import torch\n",
+    "## About the `tslib` models\n",
     "\n",
-    "from pytorch_forecasting.data.data_module import TslibDataModule\n",
-    "from pytorch_forecasting.data.encoders import (\n",
-    "    NaNLabelEncoder,\n",
-    "    TorchNormalizer,\n",
-    ")\n",
-    "from pytorch_forecasting.data.timeseries import TimeSeries\n",
-    "from pytorch_forecasting.models.timexer._timexer_v2 import TimeXer"
+    "The `tslib` models in v2 (e.g. `TimeXer`, `DLinear`) are adapted from the [Time-Series-Library (`thuml`)](https://github.com/thuml/Time-Series-Library). They are **different from the encoder–decoder models** in v2: `tslib` models are built on `TslibBaseModel` and consume the batch layout produced by `TslibDataModule`, whereas the encoder–decoder models (e.g. `TFT`) are built for `EncoderDecoderTimeSeriesDataModule`. The two families are **not interchangeable** — a `tslib` model expects the context/target tensors emitted by `TslibDataModule`, not those from the encoder–decoder datamodule.\n",
+    "\n",
+    "You can train a `tslib` model in **two ways**, and this notebook shows both:\n",
+    "\n",
+    "1. **High-level package (`pkg`) API** — hand a package class (e.g. `TimeXer_pkg_v2`) three config dicts (`model_cfg`, `datamodule_cfg`, `trainer_cfg`) and call `.fit()` / `.predict()`. It builds the datamodule, model and Lightning `Trainer` for you.\n",
+    "2. **Low-level 3-stage pipeline** — build the `TslibDataModule`, model and `Trainer` yourself. More verbose, but gives full control over the trainer, callbacks and preprocessing.\n",
+    "\n",
+    "Each section below is **self-contained**: it creates its own `TimeSeries` dataset and imports what it needs, so you can follow either one on its own. We cover the high-level API first."
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "2625ed3d",
+   "id": "25a96c0b",
    "metadata": {},
    "source": [
-    "## Construct a time series dataset\n",
+    "## Create the synthetic dataset\n",
     "\n",
-    "This step requires us to build a `TimeSeries` object for creating a time series dataset, which identifies the features from a raw time series dataset. As you can see below, we are initialising a sample time series dataset."
+    "Both sections share the same raw dataframe. We generate it with the built-in `load_toydata` helper, which returns a `pandas` DataFrame of `num_series` noisy sine-wave series with numeric, categorical, known-future and static columns."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
-   "id": "a0058487",
+   "execution_count": 1,
+   "id": "a97c1750",
    "metadata": {},
    "outputs": [
     {
@@ -93,55 +89,55 @@
        "      <th>0</th>\n",
        "      <td>0</td>\n",
        "      <td>0</td>\n",
-       "      <td>0.177658</td>\n",
-       "      <td>0.181124</td>\n",
+       "      <td>-0.100873</td>\n",
+       "      <td>0.361583</td>\n",
        "      <td>0</td>\n",
        "      <td>1.000000</td>\n",
-       "      <td>0.409581</td>\n",
+       "      <td>0.029542</td>\n",
        "      <td>0</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>0</td>\n",
        "      <td>1</td>\n",
-       "      <td>0.181124</td>\n",
-       "      <td>0.314081</td>\n",
+       "      <td>0.361583</td>\n",
+       "      <td>0.429515</td>\n",
        "      <td>0</td>\n",
        "      <td>0.995004</td>\n",
-       "      <td>0.409581</td>\n",
+       "      <td>0.029542</td>\n",
        "      <td>0</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>0</td>\n",
        "      <td>2</td>\n",
-       "      <td>0.314081</td>\n",
-       "      <td>0.601934</td>\n",
+       "      <td>0.429515</td>\n",
+       "      <td>0.672608</td>\n",
        "      <td>0</td>\n",
        "      <td>0.980067</td>\n",
-       "      <td>0.409581</td>\n",
+       "      <td>0.029542</td>\n",
        "      <td>0</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
        "      <td>0</td>\n",
        "      <td>3</td>\n",
-       "      <td>0.601934</td>\n",
-       "      <td>0.733805</td>\n",
+       "      <td>0.672608</td>\n",
+       "      <td>0.898369</td>\n",
        "      <td>0</td>\n",
        "      <td>0.955336</td>\n",
-       "      <td>0.409581</td>\n",
+       "      <td>0.029542</td>\n",
        "      <td>0</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>4</th>\n",
        "      <td>0</td>\n",
        "      <td>4</td>\n",
-       "      <td>0.733805</td>\n",
-       "      <td>0.768843</td>\n",
+       "      <td>0.898369</td>\n",
+       "      <td>1.002088</td>\n",
        "      <td>0</td>\n",
        "      <td>0.921061</td>\n",
-       "      <td>0.409581</td>\n",
+       "      <td>0.029542</td>\n",
        "      <td>0</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
@@ -150,54 +146,37 @@
       ],
       "text/plain": [
        "   series_id  time_idx         x         y  category  future_known_feature  \\\n",
-       "0          0         0  0.177658  0.181124         0              1.000000   \n",
-       "1          0         1  0.181124  0.314081         0              0.995004   \n",
-       "2          0         2  0.314081  0.601934         0              0.980067   \n",
-       "3          0         3  0.601934  0.733805         0              0.955336   \n",
-       "4          0         4  0.733805  0.768843         0              0.921061   \n",
+       "0          0         0 -0.100873  0.361583         0              1.000000   \n",
+       "1          0         1  0.361583  0.429515         0              0.995004   \n",
+       "2          0         2  0.429515  0.672608         0              0.980067   \n",
+       "3          0         3  0.672608  0.898369         0              0.955336   \n",
+       "4          0         4  0.898369  1.002088         0              0.921061   \n",
        "\n",
        "   static_feature  static_feature_cat  \n",
-       "0        0.409581                   0  \n",
-       "1        0.409581                   0  \n",
-       "2        0.409581                   0  \n",
-       "3        0.409581                   0  \n",
-       "4        0.409581                   0  "
+       "0        0.029542                   0  \n",
+       "1        0.029542                   0  \n",
+       "2        0.029542                   0  \n",
+       "3        0.029542                   0  \n",
+       "4        0.029542                   0  "
       ]
      },
-     "execution_count": 2,
+     "execution_count": 1,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "num_series = 100\n",
-    "seq_length = 50\n",
-    "data_list = []\n",
-    "for i in range(num_series):\n",
-    "    x = np.arange(seq_length)\n",
-    "    y = np.sin(x / 5.0) + np.random.normal(scale=0.1, size=seq_length)\n",
-    "    category = i % 5\n",
-    "    static_value = np.random.rand()\n",
-    "    for t in range(seq_length - 1):\n",
-    "        data_list.append(\n",
-    "            {\n",
-    "                \"series_id\": i,\n",
-    "                \"time_idx\": t,\n",
-    "                \"x\": y[t],\n",
-    "                \"y\": y[t + 1],\n",
-    "                \"category\": category,\n",
-    "                \"future_known_feature\": np.cos(t / 10),\n",
-    "                \"static_feature\": static_value,\n",
-    "                \"static_feature_cat\": i % 3,\n",
-    "            }\n",
-    "        )\n",
-    "data_df = pd.DataFrame(data_list)\n",
+    "from pytorch_forecasting.data.examples import load_toydata\n",
+    "\n",
+    "num_series = 100  # number of individual time series\n",
+    "seq_length = 50  # length of each series\n",
+    "data_df = load_toydata(num_series, seq_length)\n",
     "data_df.head()"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "c7c04ff5",
+   "id": "bd371520",
    "metadata": {},
    "source": [
     "## Feature Categories and Definitions\n",
@@ -243,17 +222,75 @@
     "- **Usage**: Entity-specific characteristics that don't change over time"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "ed836f71",
+   "metadata": {},
+   "source": [
+    "> **Note on categorical features.** The `cat` columns above are declared and carried through the `TslibDataModule` (as `history_cat` / `future_cat` / `static_categorical_features`), but **TimeXer v2 currently uses only the continuous features in the context window** — categorical variables are not yet consumed by the model (`enc_in` is set to the continuous dimension only).\n",
+    "\n",
+    "Also note the datamodule does **not** encode categoricals: every feature is cast to `float32` and sliced by index, so categorical columns must already be numeric (string categories would fail). There is no categorical-encoder step in this pipeline yet. Support for proper categorical handling is planned for a future release. We keep the `cat` declaration here to show the full data definition, not because it affects the forecast today."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b0fff5e5",
+   "metadata": {},
+   "source": [
+    "## 1. High-level package (`pkg`) API"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b7f87b84",
+   "metadata": {},
+   "source": [
+    "### Steps\n",
+    "* Create the `TimeSeries` dataset\n",
+    "* Create the config dicts (`datamodule_cfg`, `trainer_cfg`, and a per-model `model_cfg`)\n",
+    "* Create the package object (e.g. `TimeXer_pkg_v2`)\n",
+    "* Call `pkg.fit(dataset)` and `pkg.predict(dataset)`\n",
+    "\n",
+    "The package class is a thin wrapper that manages the model, its `TslibDataModule` and the Lightning `Trainer`. Checkpoints are saved automatically under `checkpoints/` unless you pass `save_ckpt=False`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1debca4a",
+   "metadata": {},
+   "source": [
+    "### Create the `TimeSeries` dataset\n",
+    "\n",
+    "`TimeSeries` turns the raw dataframe into tensors and records which columns are the target, groups, numeric/categorical, known/unknown and static features (see the reference above)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "4f0f37de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sklearn.preprocessing import StandardScaler\n",
+    "\n",
+    "from pytorch_forecasting.data.encoders import TorchNormalizer\n",
+    "from pytorch_forecasting.data.timeseries import TimeSeries\n",
+    "from pytorch_forecasting.metrics import MAE, SMAPE, QuantileLoss\n",
+    "from pytorch_forecasting.models.timexer import TimeXer_pkg_v2\n",
+    "from pytorch_forecasting.models.dlinear import DLinear_pkg_v2"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "89a5adbe",
+   "id": "7fb824fc",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/home/aryan/pytorch-forecasting/pytorch_forecasting/data/timeseries/_timeseries_v2.py:105: UserWarning: TimeSeries is part of an experimental rework of the pytorch-forecasting data layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. For beta testing, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/data/timeseries/_timeseries_v2.py:104: UserWarning: TimeSeries is part of an experimental rework of the pytorch-forecasting data layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. For beta testing, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n",
       "  warn(\n"
      ]
     }
@@ -264,7 +301,7 @@
     "    time=\"time_idx\",\n",
     "    target=\"y\",\n",
     "    group=[\"series_id\"],\n",
-    "    num=[\"x\", \"future_know_feature\", \"static_feature\"],\n",
+    "    num=[\"x\", \"future_known_feature\", \"static_feature\"],\n",
     "    cat=[\"category\", \"static_feature_cat\"],\n",
     "    known=[\"future_known_feature\"],\n",
     "    unknown=[\"x\", \"category\"],\n",
@@ -274,377 +311,281 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f8753a6a",
+   "id": "1692afdf",
    "metadata": {},
    "source": [
-    "## Initialise the `TslibDataModule` using the dataset\n",
+    "### Create the configs\n",
     "\n",
-    "This steps initialises a basic data module built specially for `tslib` modules and provides all the metadata required to train and implement the `tslib` of your choice!\n",
-    "You can refer the implementation for `TslibDataModule` for more information."
+    "`datamodule_cfg` and `trainer_cfg` are shared across the models below; each model gets its own `model_cfg` inline in the package call."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 4,
-   "id": "5eae9035",
+   "id": "d585fd06",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/aryan/pytorch-forecasting/pytorch_forecasting/data/_tslib_data_module.py:275: UserWarning: TslibDataModule is experimental and subject to change. The API is not stable and may change without prior warning.\n",
-      "  warnings.warn(\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "data_module = TslibDataModule(\n",
-    "    time_series_dataset=dataset,\n",
-    "    context_length=30,\n",
+    "datamodule_cfg = dict(\n",
+    "    context_length=32,  # divisible by patch_length(4) -> no dropped-timestep warning\n",
     "    prediction_length=1,\n",
     "    add_relative_time_idx=True,\n",
     "    target_normalizer=TorchNormalizer(),\n",
-    "    categorical_encoders={\n",
-    "        \"category\": NaNLabelEncoder(add_nan=True),\n",
-    "        \"static_feature_cat\": NaNLabelEncoder(add_nan=True),\n",
-    "    },\n",
     "    scalers={\n",
     "        \"x\": StandardScaler(),\n",
     "        \"future_known_feature\": StandardScaler(),\n",
     "        \"static_feature\": StandardScaler(),\n",
     "    },\n",
     "    batch_size=32,\n",
+    ")\n",
+    "\n",
+    "trainer_cfg = dict(\n",
+    "    max_epochs=5,\n",
+    "    accelerator=\"cpu\",\n",
+    "    devices=1,\n",
+    "    enable_progress_bar=False,\n",
+    "    enable_model_summary=False,\n",
     ")"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "b1843233",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'feature_names': {'categorical': ['category', 'static_feature_cat'],\n",
-       "  'continuous': ['x', 'future_known_feature', 'static_feature'],\n",
-       "  'static': ['static_feature', 'static_feature_cat'],\n",
-       "  'known': ['future_known_feature'],\n",
-       "  'unknown': ['x', 'category', 'static_feature', 'static_feature_cat'],\n",
-       "  'target': ['y'],\n",
-       "  'all': ['x',\n",
-       "   'category',\n",
-       "   'future_known_feature',\n",
-       "   'static_feature',\n",
-       "   'static_feature_cat'],\n",
-       "  'static_categorical': ['static_feature_cat'],\n",
-       "  'static_continuous': ['static_feature']},\n",
-       " 'feature_indices': {'categorical': [1, 4],\n",
-       "  'continuous': [0, 2, 3],\n",
-       "  'static': [],\n",
-       "  'known': [2],\n",
-       "  'unknown': [0, 1, 3, 4],\n",
-       "  'target': [0]},\n",
-       " 'n_features': {'categorical': 2,\n",
-       "  'continuous': 3,\n",
-       "  'static': 2,\n",
-       "  'known': 1,\n",
-       "  'unknown': 4,\n",
-       "  'target': 1,\n",
-       "  'all': 5,\n",
-       "  'static_categorical': 1,\n",
-       "  'static_continuous': 1},\n",
-       " 'context_length': 30,\n",
-       " 'prediction_length': 1,\n",
-       " 'freq': 'h',\n",
-       " 'features': 'MS'}"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "data_module.metadata"
-   ]
-  },
   {
    "cell_type": "markdown",
-   "id": "dd9451ee",
+   "id": "6cdfe14d",
    "metadata": {},
    "source": [
-    "## Initialise the model\n",
-    "\n",
-    "We shall try out two versions of this model, one using `MAE()` and one with `QuantileLoss()`.\n",
+    "### TimeXer with MAE (point forecast) via the package class\n",
     "\n",
-    "Let us quickly import the required packages for the next steps."
+    "Instead of wiring up the model, datamodule and `Trainer` by hand, we hand three config dicts to `TimeXer_pkg_v2` and call `.fit()` / `.predict()`."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "id": "f6b568a5",
+   "execution_count": 5,
+   "id": "48c3a097",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'loss': MAE(), 'hidden_size': 64, 'n_heads': 4, 'e_layers': 2, 'd_ff': 256, 'dropout': 0.1, 'patch_length': 4, 'logging_metrics': [MAE(), SMAPE()], 'optimizer': 'adam', 'optimizer_params': {'lr': 0.001}, 'lr_scheduler': 'reduce_lr_on_plateau', 'lr_scheduler_params': {'mode': 'min', 'factor': 0.5, 'patience': 5}}\n"
+     ]
+    }
+   ],
    "source": [
-    "from pytorch_forecasting.metrics import MAE, SMAPE, QuantileLoss"
+    "pkg_mae = TimeXer_pkg_v2(\n",
+    "    model_cfg=dict(\n",
+    "        loss=MAE(),\n",
+    "        hidden_size=64,\n",
+    "        n_heads=4,\n",
+    "        e_layers=2,\n",
+    "        d_ff=256,\n",
+    "        dropout=0.1,\n",
+    "        patch_length=4,\n",
+    "        logging_metrics=[MAE(), SMAPE()],\n",
+    "        optimizer=\"adam\",\n",
+    "        optimizer_params={\"lr\": 1e-3},\n",
+    "        lr_scheduler=\"reduce_lr_on_plateau\",\n",
+    "        lr_scheduler_params={\"mode\": \"min\", \"factor\": 0.5, \"patience\": 5},\n",
+    "    ),\n",
+    "    datamodule_cfg=datamodule_cfg,\n",
+    "    trainer_cfg=trainer_cfg,\n",
+    ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
-   "id": "429b5f15",
+   "execution_count": 6,
+   "id": "94e0f688",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/home/aryan/pytorch-forecasting/pytorch_forecasting/models/base/_base_model_v2.py:61: UserWarning: The Model 'TimeXer' is part of an experimental reworkof the pytorch-forecasting model layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. This class is intended for beta testing and as a basic skeleton, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/data/data_module/_tslib_data_module.py:331: UserWarning: TslibDataModule is experimental and subject to change. The API is not stable and may change without prior warning.\n",
+      "  warnings.warn(\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/base/_base_model_v2.py:85: UserWarning: The Model 'TimeXer' is part of an experimental reworkof the pytorch-forecasting model layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. This class is intended for beta testing and as a basic skeleton, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n",
       "  warn(\n",
-      "/home/aryan/pytorch-forecasting/pytorch_forecasting/models/base/_tslib_base_model_v2.py:60: UserWarning: The Model 'TimeXer' is part of an experimental implementationof the pytorch-forecasting model layer for Time Series Library, scheduledfor release with v2.0.0. The API is not stableand may change without prior warning. This class is intended for betatesting, not for stable production use.\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/base/_tslib_base_model_v2.py:63: UserWarning: The Model 'TimeXer' is part of an experimental implementationof the pytorch-forecasting model layer for Time Series Library, scheduledfor release with v2.0.0. The API is not stableand may change without prior warning. This class is intended for betatesting, not for stable production use.\n",
       "  warn(\n",
-      "/home/aryan/pytorch-forecasting/pytorch_forecasting/models/timexer/_timexer_v2.py:133: UserWarning: TimeXer is an experimental model implemented on TslibBaseModelV2. It is an unstable version and maybe subject to unannouced changes.Please use with caution. Feedback on the design and implementation iswelcome. On the issue #1833 - https://github.com/sktime/pytorch-forecasting/issues/1833\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/timexer/_timexer_v2.py:134: UserWarning: TimeXer is an experimental model implemented on TslibBaseModelV2. It is an unstable version and maybe subject to unannouced changes.Please use with caution. Feedback on the design and implementation iswelcome. On the issue #1833 - https://github.com/sktime/pytorch-forecasting/issues/1833\n",
       "  warn.warn(\n",
-      "/home/aryan/pytorch-forecasting/pytorch_forecasting/models/timexer/_timexer_v2.py:179: UserWarning: Context length (30) is not divisible by patch length. This may lead to unexpected behavior, as sometime steps will not be used in the model.\n",
-      "  warn.warn(\n"
+      "GPU available: True (mps), used: False\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "TPU available: False, using: 0 TPU cores\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/echoooooo/miniconda3/lib/python3.14/site-packages/lightning/pytorch/trainer/setup.py:175: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.\n",
+      ""
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/echoooooo/miniconda3/lib/python3.14/site-packages/lightning/pytorch/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.\n",
+      "/Users/echoooooo/miniconda3/lib/python3.14/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.\n",
+      "/Users/echoooooo/miniconda3/lib/python3.14/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.\n",
+      "/Users/echoooooo/miniconda3/lib/python3.14/site-packages/lightning/pytorch/loops/fit_loop.py:321: The number of training batches (38) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "`Trainer.fit` stopped: `max_epochs=5` reached.\n"
      ]
     }
    ],
    "source": [
-    "model1 = TimeXer(\n",
-    "    loss=MAE(),\n",
-    "    hidden_size=64,\n",
-    "    nhead=4,\n",
-    "    e_layers=2,\n",
-    "    d_ff=256,\n",
-    "    dropout=0.1,\n",
-    "    patch_length=4,\n",
-    "    logging_metrics=[MAE(), SMAPE()],\n",
-    "    optimizer=\"adam\",\n",
-    "    optimizer_params={\"lr\": 1e-3},\n",
-    "    lr_scheduler=\"reduce_lr_on_plateau\",\n",
-    "    lr_scheduler_params={\n",
-    "        \"mode\": \"min\",\n",
-    "        \"factor\": 0.5,\n",
-    "        \"patience\": 5,\n",
-    "    },\n",
-    "    metadata=data_module.metadata,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "id": "0aa21f48",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model2 = TimeXer(\n",
-    "    loss=QuantileLoss(quantiles=[0.1, 0.5, 0.9]),  # quantiles of 0.1, 0.5 and 0.9 used.\n",
-    "    hidden_size=64,\n",
-    "    nhead=4,\n",
-    "    e_layers=2,\n",
-    "    d_ff=256,\n",
-    "    dropout=0.1,\n",
-    "    patch_length=4,\n",
-    "    optimizer=\"adam\",\n",
-    "    optimizer_params={\"lr\": 1e-3},\n",
-    "    lr_scheduler=\"reduce_lr_on_plateau\",\n",
-    "    lr_scheduler_params={\n",
-    "        \"mode\": \"min\",\n",
-    "        \"factor\": 0.5,\n",
-    "        \"patience\": 5,\n",
-    "    },\n",
-    "    metadata=data_module.metadata,\n",
-    ")"
+    "pkg_mae.fit(dataset, save_ckpt=False)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
-   "id": "02605f9b",
+   "execution_count": 7,
+   "id": "3ace66cc",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "GPU available: True (cuda), used: True\n",
-      "TPU available: False, using: 0 TPU cores\n",
-      "HPU available: False, using: 0 HPUs\n",
-      "GPU available: True (cuda), used: True\n",
-      "TPU available: False, using: 0 TPU cores\n",
-      "HPU available: False, using: 0 HPUs\n"
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/data/data_module/_tslib_data_module.py:331: UserWarning: TslibDataModule is experimental and subject to change. The API is not stable and may change without prior warning.\n",
+      "  warnings.warn(\n",
+      "GPU available: True (mps), used: False\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "TPU available: False, using: 0 TPU cores\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/echoooooo/miniconda3/lib/python3.14/site-packages/lightning/pytorch/trainer/setup.py:175: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.\n",
+      ""
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/echoooooo/miniconda3/lib/python3.14/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'predict_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Point prediction shape: torch.Size([1700, 1])\n",
+      "Point prediction: tensor([[ 0.2879],\n",
+      "        [ 0.4023],\n",
+      "        [ 0.5390],\n",
+      "        ...,\n",
+      "        [-0.0669],\n",
+      "        [-0.2310],\n",
+      "        [-0.4071]])\n"
      ]
     }
    ],
    "source": [
-    "from lightning.pytorch import Trainer\n",
-    "\n",
-    "trainer1 = Trainer(\n",
-    "    max_epochs=5,\n",
-    "    accelerator=\"auto\",\n",
-    "    devices=1,\n",
-    "    enable_progress_bar=True,\n",
-    "    enable_model_summary=True,\n",
+    "preds_mae = pkg_mae.predict(\n",
+    "    dataset,\n",
+    "    mode=\"prediction\",\n",
+    "    trainer_kwargs={\"accelerator\": \"cpu\", \"enable_progress_bar\": False},\n",
     ")\n",
-    "\n",
-    "trainer2 = Trainer(\n",
-    "    max_epochs=4,\n",
-    "    accelerator=\"auto\",\n",
-    "    devices=1,\n",
-    "    enable_progress_bar=True,\n",
-    "    enable_model_summary=True,\n",
-    ")"
+    "print(\"Point prediction shape:\", preds_mae[\"prediction\"].shape)\n",
+    "print(\"Point prediction:\", preds_mae[\"prediction\"])"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "e22756b2",
+   "id": "099c5cb2",
    "metadata": {},
    "source": [
-    "## Fit the trainer on the model and feed data using the data module"
+    "### TimeXer with QuantileLoss (interval forecast)\n",
+    "\n",
+    "Swap the loss to `QuantileLoss` to get quantile / prediction-interval outputs."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
-   "id": "6e9117d2",
+   "execution_count": 8,
+   "id": "b355ba52",
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
-      "You are using a CUDA device ('NVIDIA GeForce RTX 4050 Laptop GPU') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
-      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
-      "\n",
-      "  | Name         | Type                   | Params | Mode \n",
-      "----------------------------------------------------------------\n",
-      "0 | loss         | MAE                    | 0      | train\n",
-      "1 | en_embedding | EnEmbedding            | 320    | train\n",
-      "2 | ex_embedding | DataEmbedding_inverted | 2.0 K  | train\n",
-      "3 | encoder      | Encoder                | 133 K  | train\n",
-      "4 | head         | FlattenHead            | 513    | train\n",
-      "----------------------------------------------------------------\n",
-      "136 K     Trainable params\n",
-      "0         Non-trainable params\n",
-      "136 K     Total params\n",
-      "0.546     Total estimated model params size (MB)\n",
-      "57        Modules in train mode\n",
-      "0         Modules in eval mode\n"
+      "{'loss': QuantileLoss(quantiles=[0.1, 0.5, 0.9]), 'hidden_size': 64, 'n_heads': 4, 'e_layers': 2, 'd_ff': 256, 'dropout': 0.1, 'patch_length': 4, 'optimizer': 'adam', 'optimizer_params': {'lr': 0.001}, 'lr_scheduler': 'reduce_lr_on_plateau', 'lr_scheduler_params': {'mode': 'min', 'factor': 0.5, 'patience': 5}}\n"
      ]
-    },
+    }
+   ],
+   "source": [
+    "pkg_q = TimeXer_pkg_v2(\n",
+    "    model_cfg=dict(\n",
+    "        loss=QuantileLoss(quantiles=[0.1, 0.5, 0.9]),\n",
+    "        hidden_size=64,\n",
+    "        n_heads=4,\n",
+    "        e_layers=2,\n",
+    "        d_ff=256,\n",
+    "        dropout=0.1,\n",
+    "        patch_length=4,\n",
+    "        # NOTE: no logging_metrics here — point metrics (MAE/SMAPE) cannot be\n",
+    "        # computed on QuantileLoss output (extra quantile dim) and would raise\n",
+    "        # \"Prediction should only have one extra dimension\" during fit.\n",
+    "        optimizer=\"adam\",\n",
+    "        optimizer_params={\"lr\": 1e-3},\n",
+    "        lr_scheduler=\"reduce_lr_on_plateau\",\n",
+    "        lr_scheduler_params={\"mode\": \"min\", \"factor\": 0.5, \"patience\": 5},\n",
+    "    ),\n",
+    "    datamodule_cfg=datamodule_cfg,\n",
+    "    trainer_cfg=trainer_cfg,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "e20c7e0a",
+   "metadata": {},
+   "outputs": [
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f26d868819404cb0a48cc030aefef48c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/base/_base_model_v2.py:85: UserWarning: The Model 'TimeXer' is part of an experimental reworkof the pytorch-forecasting model layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. This class is intended for beta testing and as a basic skeleton, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n",
+      "  warn(\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/base/_tslib_base_model_v2.py:63: UserWarning: The Model 'TimeXer' is part of an experimental implementationof the pytorch-forecasting model layer for Time Series Library, scheduledfor release with v2.0.0. The API is not stableand may change without prior warning. This class is intended for betatesting, not for stable production use.\n",
+      "  warn(\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/timexer/_timexer_v2.py:134: UserWarning: TimeXer is an experimental model implemented on TslibBaseModelV2. It is an unstable version and maybe subject to unannouced changes.Please use with caution. Feedback on the design and implementation iswelcome. On the issue #1833 - https://github.com/sktime/pytorch-forecasting/issues/1833\n",
+      "  warn.warn(\n",
+      "GPU available: True (mps), used: False\n"
+     ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/home/aryan/pytorch-forecasting/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:425: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=15` in the `DataLoader` to improve performance.\n",
-      "/home/aryan/pytorch-forecasting/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:425: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=15` in the `DataLoader` to improve performance.\n",
-      "/home/aryan/pytorch-forecasting/.venv/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py:310: The number of training batches (42) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n"
+      "TPU available: False, using: 0 TPU cores\n"
      ]
     },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "db0fdcf01b784ddf8bf0d93fe1cc23e4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Training: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "26073d89e3014850ba0c4d7f01104047",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bf42294210dc4d53a066967bc54b1282",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "27ed3f4556284bf0a918d22554e8e9c2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0e338f265f554e6d96020135a3f64f64",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "46bbd28e869740449aad2c7c1aaf5da7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
     {
      "name": "stderr",
      "output_type": "stream",
@@ -654,602 +595,1387 @@
     }
    ],
    "source": [
-    "trainer1.fit(model1, data_module)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e8fb4f31",
-   "metadata": {},
-   "source": [
-    "Now let us train the model using `QuantileLoss`."
+    "pkg_q.fit(dataset, save_ckpt=False)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
-   "id": "3c67d86f",
+   "execution_count": 10,
+   "id": "53c8c3b8",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
-      "\n",
-      "  | Name         | Type                   | Params | Mode \n",
-      "----------------------------------------------------------------\n",
-      "0 | loss         | QuantileLoss           | 0      | train\n",
-      "1 | en_embedding | EnEmbedding            | 320    | train\n",
-      "2 | ex_embedding | DataEmbedding_inverted | 2.0 K  | train\n",
-      "3 | encoder      | Encoder                | 133 K  | train\n",
-      "4 | head         | FlattenHead            | 1.5 K  | train\n",
-      "----------------------------------------------------------------\n",
-      "137 K     Trainable params\n",
-      "0         Non-trainable params\n",
-      "137 K     Total params\n",
-      "0.550     Total estimated model params size (MB)\n",
-      "57        Modules in train mode\n",
-      "0         Modules in eval mode\n"
+      "GPU available: True (mps), used: False\n"
      ]
     },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "acba85c2e99b4dc391315e10471e1e8c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/home/aryan/pytorch-forecasting/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:425: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=15` in the `DataLoader` to improve performance.\n",
-      "/home/aryan/pytorch-forecasting/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:425: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=15` in the `DataLoader` to improve performance.\n",
-      "/home/aryan/pytorch-forecasting/.venv/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py:310: The number of training batches (42) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n"
+      "TPU available: False, using: 0 TPU cores\n"
      ]
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "77d8d477ad3040aea894b76ca2df8130",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Training: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f1bbc1f476b046e788d827e5dc02d888",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1233ff4c4fbe4bde9e7717c54e28567d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "334cb342267b47b1b4e1e2ca1038569f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "26260a641c604444811ec1bf00c80edb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
-      "`Trainer.fit` stopped: `max_epochs=4` reached.\n"
+      "Quantile prediction shape (batch, pred_len, n_quantiles): torch.Size([1700, 1, 3])\n",
+      "Quantile prediction: tensor([[[ 0.1936,  0.4580,  0.5065]],\n",
+      "\n",
+      "        [[ 0.3599,  0.6098,  0.6686]],\n",
+      "\n",
+      "        [[ 0.4852,  0.7296,  0.7853]],\n",
+      "\n",
+      "        ...,\n",
+      "\n",
+      "        [[-0.2017,  0.1089,  0.2374]],\n",
+      "\n",
+      "        [[-0.3642, -0.0622,  0.0831]],\n",
+      "\n",
+      "        [[-0.5479, -0.2107, -0.0850]]])\n"
      ]
     }
    ],
    "source": [
-    "trainer2.fit(model2, data_module)"
+    "preds_q = pkg_q.predict(\n",
+    "    dataset,\n",
+    "    mode=\"quantiles\",\n",
+    "    trainer_kwargs={\"accelerator\": \"cpu\", \"enable_progress_bar\": False},\n",
+    ")\n",
+    "print(\n",
+    "    \"Quantile prediction shape (batch, pred_len, n_quantiles):\",\n",
+    "    preds_q[\"prediction\"].shape,\n",
+    ")\n",
+    "print(\"Quantile prediction:\", preds_q[\"prediction\"])"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "16e2d445",
+   "id": "831949eb",
    "metadata": {},
    "source": [
-    "## Test the model"
+    "### Same pipeline, different model: DLinear\n",
+    "\n",
+    "`DLinear` is the other model built on the `tslib` base, so it runs through the exact same pipeline — only the package class and its `model_cfg` change."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
-   "id": "dbf1ace6",
+   "execution_count": 11,
+   "id": "e3f0201a",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
-      "/home/aryan/pytorch-forecasting/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:425: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=15` in the `DataLoader` to improve performance.\n"
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/base/_base_model_v2.py:85: UserWarning: The Model 'DLinear' is part of an experimental reworkof the pytorch-forecasting model layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. This class is intended for beta testing and as a basic skeleton, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n",
+      "  warn(\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/base/_tslib_base_model_v2.py:63: UserWarning: The Model 'DLinear' is part of an experimental implementationof the pytorch-forecasting model layer for Time Series Library, scheduledfor release with v2.0.0. The API is not stableand may change without prior warning. This class is intended for betatesting, not for stable production use.\n",
+      "  warn(\n",
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/models/dlinear/_dlinear_v2.py:96: UserWarning: DLinear is an experimental model implemented on TslibBaseModelV2. It is an unstable version and may be subject to unannounced changes. Please use with caution.\n",
+      "  warnings.warn(\n",
+      "GPU available: True (mps), used: False\n"
      ]
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d8e51a5bb1564edab4e0d32743a16299",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Testing: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "TPU available: False, using: 0 TPU cores\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'loss': MAE(), 'moving_avg': 25, 'individual': False, 'logging_metrics': [MAE(), SMAPE()], 'optimizer': 'adam', 'optimizer_params': {'lr': 0.001}}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "`Trainer.fit` stopped: `max_epochs=5` reached.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (mps), used: False\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "TPU available: False, using: 0 TPU cores\n"
+     ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
-      "       Test metric             DataLoader 0\n",
-      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
-      "        test_MAE            0.13894660770893097\n",
-      "       test_SMAPE           0.40041154623031616\n",
-      "        test_loss           0.13894660770893097\n",
-      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
+      "DLinear prediction shape: torch.Size([1700, 1])\n",
+      "DLinear prediction: tensor([[ 0.3987],\n",
+      "        [ 0.5763],\n",
+      "        [ 0.7305],\n",
+      "        ...,\n",
+      "        [ 0.0261],\n",
+      "        [-0.1774],\n",
+      "        [-0.3777]])\n"
      ]
     }
    ],
    "source": [
-    "test_metrics = trainer1.test(model1, data_module)"
+    "pkg_dl = DLinear_pkg_v2(\n",
+    "    model_cfg=dict(\n",
+    "        loss=MAE(),\n",
+    "        moving_avg=25,\n",
+    "        individual=False,\n",
+    "        logging_metrics=[MAE(), SMAPE()],\n",
+    "        optimizer=\"adam\",\n",
+    "        optimizer_params={\"lr\": 1e-3},\n",
+    "    ),\n",
+    "    datamodule_cfg=datamodule_cfg,\n",
+    "    trainer_cfg=trainer_cfg,\n",
+    ")\n",
+    "pkg_dl.fit(dataset, save_ckpt=False)\n",
+    "preds_dl = pkg_dl.predict(\n",
+    "    dataset,\n",
+    "    mode=\"prediction\",\n",
+    "    trainer_kwargs={\"accelerator\": \"cpu\", \"enable_progress_bar\": False},\n",
+    ")\n",
+    "print(\"DLinear prediction shape:\", preds_dl[\"prediction\"].shape)\n",
+    "print(\"DLinear prediction:\", preds_dl[\"prediction\"])"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 13,
-   "id": "250b128a",
+   "cell_type": "markdown",
+   "id": "0fbf38fc",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "TimeXer(\n",
-       "  (loss): MAE()\n",
-       "  (en_embedding): EnEmbedding(\n",
-       "    (value_embedding): Linear(in_features=4, out_features=64, bias=False)\n",
-       "    (position_embedding): PositionalEmbedding()\n",
-       "    (dropout): Dropout(p=0.1, inplace=False)\n",
-       "  )\n",
-       "  (ex_embedding): DataEmbedding_inverted(\n",
-       "    (value_embedding): Linear(in_features=30, out_features=64, bias=True)\n",
-       "    (dropout): Dropout(p=0.1, inplace=False)\n",
-       "  )\n",
-       "  (encoder): Encoder(\n",
-       "    (layers): ModuleList(\n",
-       "      (0-1): 2 x EncoderLayer(\n",
-       "        (self_attention): AttentionLayer(\n",
-       "          (inner_attention): FullAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "          )\n",
-       "          (query_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (key_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (value_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (out_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "        )\n",
-       "        (cross_attention): AttentionLayer(\n",
-       "          (inner_attention): FullAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "          )\n",
-       "          (query_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (key_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (value_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (out_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "        )\n",
-       "        (conv1): Conv1d(64, 256, kernel_size=(1,), stride=(1,))\n",
-       "        (conv2): Conv1d(256, 64, kernel_size=(1,), stride=(1,))\n",
-       "        (norm1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
-       "        (norm2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
-       "        (norm3): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
-       "  )\n",
-       "  (head): FlattenHead(\n",
-       "    (flatten): Flatten(start_dim=-2, end_dim=-1)\n",
-       "    (linear): Linear(in_features=512, out_features=1, bias=True)\n",
-       "    (dropout): Dropout(p=0.1, inplace=False)\n",
-       "  )\n",
-       ")"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
    "source": [
-    "model1.eval()"
+    "## 2. Low-level 3-stage pipeline"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 14,
-   "id": "f730b49a",
+   "cell_type": "markdown",
+   "id": "11d1710e",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Prediction: tensor([[[ 0.1253]],\n",
-      "\n",
-      "        [[ 0.2623]],\n",
-      "\n",
-      "        [[ 0.4591]],\n",
-      "\n",
-      "        [[ 0.6304]],\n",
-      "\n",
-      "        [[ 0.7916]],\n",
-      "\n",
-      "        [[ 0.9132]],\n",
-      "\n",
-      "        [[ 1.0252]],\n",
-      "\n",
-      "        [[ 1.1069]],\n",
-      "\n",
-      "        [[ 1.1370]],\n",
-      "\n",
-      "        [[ 1.1317]],\n",
-      "\n",
-      "        [[ 1.0659]],\n",
-      "\n",
-      "        [[ 0.9617]],\n",
-      "\n",
-      "        [[ 0.8297]],\n",
-      "\n",
-      "        [[ 0.6622]],\n",
-      "\n",
-      "        [[ 0.5254]],\n",
-      "\n",
-      "        [[ 0.3310]],\n",
-      "\n",
-      "        [[ 0.1579]],\n",
-      "\n",
-      "        [[-0.0506]],\n",
-      "\n",
-      "        [[-0.1999]],\n",
-      "\n",
-      "        [[ 0.0740]],\n",
-      "\n",
-      "        [[ 0.2787]],\n",
-      "\n",
-      "        [[ 0.4506]],\n",
-      "\n",
-      "        [[ 0.6381]],\n",
-      "\n",
-      "        [[ 0.7867]],\n",
-      "\n",
-      "        [[ 0.9343]],\n",
-      "\n",
-      "        [[ 1.0370]],\n",
-      "\n",
-      "        [[ 1.1286]],\n",
-      "\n",
-      "        [[ 1.1737]],\n",
-      "\n",
-      "        [[ 1.1367]],\n",
-      "\n",
-      "        [[ 1.0765]],\n",
-      "\n",
-      "        [[ 0.9569]],\n",
-      "\n",
-      "        [[ 0.8583]]])\n"
-     ]
-    }
-   ],
    "source": [
-    "with torch.no_grad():\n",
-    "    test_batch = next(iter(data_module.test_dataloader()))\n",
-    "    x_test, y_test = test_batch\n",
-    "    y_pred = model1(x_test)\n",
+    "### Steps\n",
+    "1. Create the `TimeSeries` dataset\n",
+    "2. Create the `TslibDataModule`\n",
+    "3. Initialise, train and run inference with the model\n",
     "\n",
-    "    print(\"Prediction:\", y_pred[\"prediction\"])"
+    "This section does by hand what the package class did for you above. It is fully self-contained — it rebuilds the `TimeSeries` dataset from `data_df` and re-imports what it needs — so you can follow it without running section 1."
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 15,
-   "id": "e316c047",
+   "cell_type": "markdown",
+   "id": "5bfbdd89",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "torch.Size([32, 1, 1])"
-      ]
-     },
-     "execution_count": 15,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
    "source": [
-    "y_pred[\"prediction\"].shape"
+    "### 1. Create the dataset\n",
+    "\n",
+    "Same `TimeSeries` construction as above; the `data_module` will consume this and handle the dataloaders and preprocessing."
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "a01927d4",
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "988ff219",
    "metadata": {},
+   "outputs": [],
    "source": [
-    "Let us do the same for `QuantileLoss` predictions."
+    "import torch\n",
+    "from sklearn.preprocessing import StandardScaler\n",
+    "\n",
+    "from pytorch_forecasting.data.data_module import TslibDataModule\n",
+    "from pytorch_forecasting.data.encoders import TorchNormalizer\n",
+    "from pytorch_forecasting.data.timeseries import TimeSeries\n",
+    "from pytorch_forecasting.metrics import MAE, SMAPE, QuantileLoss\n",
+    "from pytorch_forecasting.models.timexer._timexer_v2 import TimeXer"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
-   "id": "22bd191f",
+   "execution_count": 13,
+   "id": "32e01d37",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
-      "/home/aryan/pytorch-forecasting/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:425: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=15` in the `DataLoader` to improve performance.\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f29e92e8bae645c0bc6692d18c8bdd1f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Testing: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
-      "       Test metric             DataLoader 0\n",
-      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
-      "        test_loss           0.07047828286886215\n",
-      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
+      "/Users/echoooooo/Desktop/code/pytorch-forecasting/pytorch_forecasting/data/timeseries/_timeseries_v2.py:104: UserWarning: TimeSeries is part of an experimental rework of the pytorch-forecasting data layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. For beta testing, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n",
+      "  warn(\n"
      ]
     }
    ],
    "source": [
-    "test_metrics = trainer2.test(model2, data_module)"
+    "dataset = TimeSeries(\n",
+    "    data=data_df,\n",
+    "    time=\"time_idx\",\n",
+    "    target=\"y\",\n",
+    "    group=[\"series_id\"],\n",
+    "    num=[\"x\", \"future_known_feature\", \"static_feature\"],\n",
+    "    cat=[\"category\", \"static_feature_cat\"],\n",
+    "    known=[\"future_known_feature\"],\n",
+    "    unknown=[\"x\", \"category\"],\n",
+    "    static=[\"static_feature\", \"static_feature_cat\"],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b641fb61",
+   "metadata": {},
+   "source": [
+    "### 2. Create the datamodule\n",
+    "\n",
+    "We build the `TslibDataModule` explicitly. We do not pass `categorical_encoders`: the `tslib` datamodule casts every feature to `float32` and does not run a categorical-encoder step (see the note on categorical features above). `setup(stage=\"fit\")` prepares the train/val splits and exposes `metadata`, which the model needs for its input/output shapes."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
-   "id": "a1d857db",
+   "execution_count": 14,
+   "id": "50915926",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "TimeXer(\n",
-       "  (loss): QuantileLoss(quantiles=[0.1, 0.5, 0.9])\n",
-       "  (en_embedding): EnEmbedding(\n",
-       "    (value_embedding): Linear(in_features=4, out_features=64, bias=False)\n",
-       "    (position_embedding): PositionalEmbedding()\n",
-       "    (dropout): Dropout(p=0.1, inplace=False)\n",
-       "  )\n",
-       "  (ex_embedding): DataEmbedding_inverted(\n",
-       "    (value_embedding): Linear(in_features=30, out_features=64, bias=True)\n",
-       "    (dropout): Dropout(p=0.1, inplace=False)\n",
-       "  )\n",
-       "  (encoder): Encoder(\n",
-       "    (layers): ModuleList(\n",
-       "      (0-1): 2 x EncoderLayer(\n",
-       "        (self_attention): AttentionLayer(\n",
-       "          (inner_attention): FullAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "          )\n",
-       "          (query_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (key_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (value_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (out_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "        )\n",
-       "        (cross_attention): AttentionLayer(\n",
-       "          (inner_attention): FullAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "          )\n",
-       "          (query_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (key_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (value_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "          (out_projection): Linear(in_features=64, out_features=64, bias=True)\n",
-       "        )\n",
-       "        (conv1): Conv1d(64, 256, kernel_size=(1,), stride=(1,))\n",
-       "        (conv2): Conv1d(256, 64, kernel_size=(1,), stride=(1,))\n",
-       "        (norm1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
-       "        (norm2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
-       "        (norm3): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
-       "  )\n",
-       "  (head): FlattenHead(\n",
-       "    (flatten): Flatten(start_dim=-2, end_dim=-1)\n",
-       "    (linear): Linear(in_features=512, out_features=3, bias=True)\n",
-       "    (dropout): Dropout(p=0.1, inplace=False)\n",
-       "  )\n",
-       ")"
-      ]
-     },
-     "execution_count": 17,
+       "{'feature_names': {'categorical': ['category', 'static_feature_cat'],\n",
+       "  'continuous': ['x', 'future_known_feature', 'static_feature'],\n",
+       "  'static': ['static_feature', 'static_feature_cat'],\n",
+       "  'known': ['future_known_feature'],\n",
+       "  'unknown': ['x', 'category', 'static_feature', 'static_feature_cat'],\n",
+       "  'target': ['y'],\n",
+       "  'all': ['x',\n",
+       "   'category',\n",
+       "   'future_known_feature',\n",
+       "   'static_feature',\n",
+       "   'static_feature_cat'],\n",
+       "  'static_categorical': ['static_feature_cat'],\n",
+       "  'static_continuous': ['static_feature']},\n",
+       " 'feature_indices': {'categorical': [1, 4],\n",
+       "  'continuous': [0, 2, 3],\n",
+       "  'static': [],\n",
+       "  'known': [2],\n",
+       "  'unknown': [0, 1, 3, 4],\n",
+       "  'target': [0]},\n",
+       " 'n_features': {'categorical': 2,\n",
+       "  'continuous': 3,\n",
+       "  'static': 2,\n",
+       "  'known': 1,\n",
+       "  'unknown': 4,\n",
+       "  'target': 1,\n",
+       "  'all': 5,\n",
+       "  'static_categorical': 1,\n",
+       "  'static_continuous': 1},\n",
+       " 'context_length': 32,\n",
+       " 'prediction_length': 1,\n",
+       " 'freq': 'h',\n",
+       " 'features': 'MS'}"
+      ]
+     },
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "model2.eval()"
+    "data_module = TslibDataModule(\n",
+    "    time_series_dataset=dataset,\n",
+    "    context_length=32,  # divisible by patch_length(4) -> no dropped-timestep warning\n",
+    "    prediction_length=1,\n",
+    "    add_relative_time_idx=True,\n",
+    "    target_normalizer=TorchNormalizer(),\n",
+    "    scalers={\n",
+    "        \"x\": StandardScaler(),\n",
+    "        \"future_known_feature\": StandardScaler(),\n",
+    "        \"static_feature\": StandardScaler(),\n",
+    "    },\n",
+    "    batch_size=32,\n",
+    ")\n",
+    "data_module.setup(stage=\"fit\")\n",
+    "data_module.metadata"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1c6d5a44",
+   "metadata": {},
+   "source": [
+    "### 3. Initialise and train the model\n",
+    "\n",
+    "We build two `TimeXer` models — one with `MAE()` (point forecast) and one with `QuantileLoss()` (interval forecast) — passing `metadata=data_module.metadata` so each model knows its input/output shapes. The `QuantileLoss` model omits `logging_metrics`: point metrics (MAE/SMAPE) cannot be computed on quantile output."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "e480232d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_mae = TimeXer(\n",
+    "    loss=MAE(),\n",
+    "    hidden_size=64,\n",
+    "    n_heads=4,\n",
+    "    e_layers=2,\n",
+    "    d_ff=256,\n",
+    "    dropout=0.1,\n",
+    "    patch_length=4,\n",
+    "    logging_metrics=[MAE(), SMAPE()],\n",
+    "    optimizer=\"adam\",\n",
+    "    optimizer_params={\"lr\": 1e-3},\n",
+    "    lr_scheduler=\"reduce_lr_on_plateau\",\n",
+    "    lr_scheduler_params={\"mode\": \"min\", \"factor\": 0.5, \"patience\": 5},\n",
+    "    metadata=data_module.metadata,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6f3bc535",
+   "metadata": {},
+   "source": [
+    "Inspect the model architecture with a plain-text `ModelSummary` (we disabled the trainer's automatic summary above to keep the training logs clean):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "25e0de8f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "  | Name            | Type                   | Params | Mode  | FLOPs\n",
+      "---------------------------------------------------------------------------\n",
+      "0 | loss            | MAE                    | 0      | train | 0    \n",
+      "1 | logging_metrics | ModuleList             | 0      | train | 0    \n",
+      "2 | en_embedding    | EnEmbedding            | 320    | train | 0    \n",
+      "3 | ex_embedding    | DataEmbedding_inverted | 2.1 K  | train | 0    \n",
+      "4 | encoder         | Encoder                | 133 K  | train | 0    \n",
+      "5 | head            | FlattenHead            | 577    | train | 0    \n",
+      "---------------------------------------------------------------------------\n",
+      "136 K     Trainable params\n",
+      "0         Non-trainable params\n",
+      "136 K     Total params\n",
+      "0.547     Total estimated model params size (MB)\n",
+      "60        Modules in train mode\n",
+      "0         Modules in eval mode\n",
+      "0         Total Flops\n"
+     ]
+    }
+   ],
+   "source": [
+    "from lightning.pytorch.utilities.model_summary import ModelSummary\n",
+    "\n",
+    "print(ModelSummary(model_mae, max_depth=1))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "f052068b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_q = TimeXer(\n",
+    "    loss=QuantileLoss(quantiles=[0.1, 0.5, 0.9]),\n",
+    "    hidden_size=64,\n",
+    "    n_heads=4,\n",
+    "    e_layers=2,\n",
+    "    d_ff=256,\n",
+    "    dropout=0.1,\n",
+    "    patch_length=4,\n",
+    "    # no logging_metrics: point metrics can't be computed on QuantileLoss output\n",
+    "    optimizer=\"adam\",\n",
+    "    optimizer_params={\"lr\": 1e-3},\n",
+    "    lr_scheduler=\"reduce_lr_on_plateau\",\n",
+    "    lr_scheduler_params={\"mode\": \"min\", \"factor\": 0.5, \"patience\": 5},\n",
+    "    metadata=data_module.metadata,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2a49bcc4",
+   "metadata": {},
+   "source": [
+    "We use a PyTorch Lightning `Trainer` to train each model. With the low-level API you create the `Trainer` yourself and call `trainer.fit(model, data_module)`."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 18,
-   "id": "52e2a36a",
+   "id": "dac60055",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (mps), used: False\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "TPU available: False, using: 0 TPU cores\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (mps), used: False\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "TPU available: False, using: 0 TPU cores\n"
+     ]
+    }
+   ],
+   "source": [
+    "from lightning.pytorch import Trainer\n",
+    "\n",
+    "trainer_mae = Trainer(\n",
+    "    max_epochs=5,\n",
+    "    accelerator=\"cpu\",\n",
+    "    devices=1,\n",
+    "    enable_progress_bar=False,\n",
+    "    enable_model_summary=False,\n",
+    ")\n",
+    "trainer_q = Trainer(\n",
+    "    max_epochs=5,\n",
+    "    accelerator=\"cpu\",\n",
+    "    devices=1,\n",
+    "    enable_progress_bar=False,\n",
+    "    enable_model_summary=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "c06b9943",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "`Trainer.fit` stopped: `max_epochs=5` reached.\n"
+     ]
+    }
+   ],
+   "source": [
+    "trainer_mae.fit(model_mae, data_module)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "711f2f1a",
+   "metadata": {},
+   "source": [
+    "Now train the `QuantileLoss` model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "e0c7e151",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "`Trainer.fit` stopped: `max_epochs=5` reached.\n"
+     ]
+    }
+   ],
+   "source": [
+    "trainer_q.fit(model_q, data_module)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d70ac0dd",
+   "metadata": {},
+   "source": [
+    "#### Output\n",
+    "\n",
+    "For inference we set up the test split and call `model.predict(test_dataloader, mode=...)`, which runs a Lightning `Trainer` internally and returns a dict with a `prediction` key. Use `mode=\"prediction\"` for the point model and `mode=\"quantiles\"` for the quantile model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "d947b3a5",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (mps), used: False\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "TPU available: False, using: 0 TPU cores\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Prediction: tensor([[[-0.1025, -0.0489,  0.0900]],\n",
+      "Point prediction shape: torch.Size([255, 1])\n",
+      "Point prediction: tensor([[ 0.3488],\n",
+      "        [ 0.5059],\n",
+      "        [ 0.6343],\n",
+      "        [ 0.7578],\n",
+      "        [ 0.8611],\n",
+      "        [ 0.9248],\n",
+      "        [ 0.9299],\n",
+      "        [ 0.9100],\n",
+      "        [ 0.8661],\n",
+      "        [ 0.7891],\n",
+      "        [ 0.6814],\n",
+      "        [ 0.5535],\n",
+      "        [ 0.3827],\n",
+      "        [ 0.2044],\n",
+      "        [ 0.0192],\n",
+      "        [-0.1669],\n",
+      "        [-0.3161],\n",
+      "        [ 0.3698],\n",
+      "        [ 0.5284],\n",
+      "        [ 0.6664],\n",
+      "        [ 0.7810],\n",
+      "        [ 0.8704],\n",
+      "        [ 0.9094],\n",
+      "        [ 0.9174],\n",
+      "        [ 0.9129],\n",
+      "        [ 0.8754],\n",
+      "        [ 0.7907],\n",
+      "        [ 0.6852],\n",
+      "        [ 0.5433],\n",
+      "        [ 0.3623],\n",
+      "        [ 0.1919],\n",
+      "        [ 0.0011],\n",
+      "        [-0.1583],\n",
+      "        [-0.3293],\n",
+      "        [ 0.3221],\n",
+      "        [ 0.4982],\n",
+      "        [ 0.6518],\n",
+      "        [ 0.7669],\n",
+      "        [ 0.8737],\n",
+      "        [ 0.9386],\n",
+      "        [ 0.9625],\n",
+      "        [ 0.9699],\n",
+      "        [ 0.9225],\n",
+      "        [ 0.8329],\n",
+      "        [ 0.7288],\n",
+      "        [ 0.5883],\n",
+      "        [ 0.4025],\n",
+      "        [ 0.2260],\n",
+      "        [ 0.0270],\n",
+      "        [-0.1689],\n",
+      "        [-0.3510],\n",
+      "        [ 0.3709],\n",
+      "        [ 0.5245],\n",
+      "        [ 0.6626],\n",
+      "        [ 0.7782],\n",
+      "        [ 0.8707],\n",
+      "        [ 0.9114],\n",
+      "        [ 0.9378],\n",
+      "        [ 0.9064],\n",
+      "        [ 0.8537],\n",
+      "        [ 0.7720],\n",
+      "        [ 0.6473],\n",
+      "        [ 0.5014],\n",
+      "        [ 0.3294],\n",
+      "        [ 0.1255],\n",
+      "        [-0.0468],\n",
+      "        [-0.2053],\n",
+      "        [-0.3445],\n",
+      "        [ 0.3461],\n",
+      "        [ 0.5025],\n",
+      "        [ 0.6340],\n",
+      "        [ 0.7541],\n",
+      "        [ 0.8552],\n",
+      "        [ 0.8974],\n",
+      "        [ 0.9192],\n",
+      "        [ 0.9146],\n",
+      "        [ 0.8652],\n",
+      "        [ 0.7836],\n",
+      "        [ 0.6785],\n",
+      "        [ 0.5342],\n",
+      "        [ 0.3724],\n",
+      "        [ 0.1925],\n",
+      "        [ 0.0052],\n",
+      "        [-0.1656],\n",
+      "        [-0.3205],\n",
+      "        [ 0.3227],\n",
+      "        [ 0.4775],\n",
+      "        [ 0.6339],\n",
+      "        [ 0.7542],\n",
+      "        [ 0.8405],\n",
+      "        [ 0.9109],\n",
+      "        [ 0.9486],\n",
+      "        [ 0.9309],\n",
+      "        [ 0.8976],\n",
+      "        [ 0.8153],\n",
+      "        [ 0.7016],\n",
+      "        [ 0.5509],\n",
+      "        [ 0.3996],\n",
+      "        [ 0.2076],\n",
+      "        [ 0.0354],\n",
+      "        [-0.1221],\n",
+      "        [-0.3016],\n",
+      "        [ 0.3404],\n",
+      "        [ 0.4948],\n",
+      "        [ 0.6283],\n",
+      "        [ 0.7628],\n",
+      "        [ 0.8589],\n",
+      "        [ 0.9118],\n",
+      "        [ 0.9261],\n",
+      "        [ 0.9265],\n",
+      "        [ 0.8945],\n",
+      "        [ 0.8029],\n",
+      "        [ 0.6820],\n",
+      "        [ 0.5675],\n",
+      "        [ 0.3966],\n",
+      "        [ 0.1972],\n",
+      "        [ 0.0257],\n",
+      "        [-0.1518],\n",
+      "        [-0.3015],\n",
+      "        [ 0.3244],\n",
+      "        [ 0.4852],\n",
+      "        [ 0.6275],\n",
+      "        [ 0.7462],\n",
+      "        [ 0.8522],\n",
+      "        [ 0.9115],\n",
+      "        [ 0.9424],\n",
+      "        [ 0.9502],\n",
+      "        [ 0.9168],\n",
+      "        [ 0.8320],\n",
+      "        [ 0.7119],\n",
+      "        [ 0.5590],\n",
+      "        [ 0.4020],\n",
+      "        [ 0.2250],\n",
+      "        [ 0.0349],\n",
+      "        [-0.1281],\n",
+      "        [-0.2896],\n",
+      "        [ 0.3270],\n",
+      "        [ 0.4817],\n",
+      "        [ 0.6173],\n",
+      "        [ 0.7281],\n",
+      "        [ 0.8225],\n",
+      "        [ 0.8909],\n",
+      "        [ 0.9342],\n",
+      "        [ 0.9350],\n",
+      "        [ 0.8892],\n",
+      "        [ 0.8050],\n",
+      "        [ 0.6944],\n",
+      "        [ 0.5455],\n",
+      "        [ 0.3916],\n",
+      "        [ 0.2246],\n",
+      "        [ 0.0224],\n",
+      "        [-0.1533],\n",
+      "        [-0.3149],\n",
+      "        [ 0.3825],\n",
+      "        [ 0.5227],\n",
+      "        [ 0.6552],\n",
+      "        [ 0.7667],\n",
+      "        [ 0.8442],\n",
+      "        [ 0.9101],\n",
+      "        [ 0.9470],\n",
+      "        [ 0.9203],\n",
+      "        [ 0.8463],\n",
+      "        [ 0.7771],\n",
+      "        [ 0.6784],\n",
+      "        [ 0.5306],\n",
+      "        [ 0.3680],\n",
+      "        [ 0.1760],\n",
+      "        [-0.0184],\n",
+      "        [-0.1717],\n",
+      "        [-0.3260],\n",
+      "        [ 0.3328],\n",
+      "        [ 0.5026],\n",
+      "        [ 0.6456],\n",
+      "        [ 0.7637],\n",
+      "        [ 0.8734],\n",
+      "        [ 0.9205],\n",
+      "        [ 0.9507],\n",
+      "        [ 0.9297],\n",
+      "        [ 0.8943],\n",
+      "        [ 0.7927],\n",
+      "        [ 0.6658],\n",
+      "        [ 0.5326],\n",
+      "        [ 0.3568],\n",
+      "        [ 0.1712],\n",
+      "        [-0.0269],\n",
+      "        [-0.2027],\n",
+      "        [-0.3453],\n",
+      "        [ 0.3091],\n",
+      "        [ 0.4582],\n",
+      "        [ 0.5828],\n",
+      "        [ 0.6981],\n",
+      "        [ 0.8046],\n",
+      "        [ 0.8765],\n",
+      "        [ 0.9052],\n",
+      "        [ 0.9249],\n",
+      "        [ 0.8764],\n",
+      "        [ 0.8027],\n",
+      "        [ 0.7078],\n",
+      "        [ 0.5764],\n",
+      "        [ 0.4075],\n",
+      "        [ 0.2525],\n",
+      "        [ 0.0670],\n",
+      "        [-0.1189],\n",
+      "        [-0.2594],\n",
+      "        [ 0.3604],\n",
+      "        [ 0.5369],\n",
+      "        [ 0.6708],\n",
+      "        [ 0.7692],\n",
+      "        [ 0.8614],\n",
+      "        [ 0.9426],\n",
+      "        [ 0.9516],\n",
+      "        [ 0.9290],\n",
+      "        [ 0.8794],\n",
+      "        [ 0.7970],\n",
+      "        [ 0.6993],\n",
+      "        [ 0.5636],\n",
+      "        [ 0.3785],\n",
+      "        [ 0.1931],\n",
+      "        [ 0.0211],\n",
+      "        [-0.1545],\n",
+      "        [-0.3122],\n",
+      "        [ 0.3462],\n",
+      "        [ 0.5141],\n",
+      "        [ 0.6610],\n",
+      "        [ 0.7852],\n",
+      "        [ 0.8467],\n",
+      "        [ 0.8948],\n",
+      "        [ 0.9374],\n",
+      "        [ 0.9398],\n",
+      "        [ 0.8970],\n",
+      "        [ 0.8152],\n",
+      "        [ 0.6901],\n",
+      "        [ 0.5336],\n",
+      "        [ 0.3569],\n",
+      "        [ 0.1655],\n",
+      "        [-0.0151],\n",
+      "        [-0.1643],\n",
+      "        [-0.3127],\n",
+      "        [ 0.3611],\n",
+      "        [ 0.5123],\n",
+      "        [ 0.6283],\n",
+      "        [ 0.7385],\n",
+      "        [ 0.8354],\n",
+      "        [ 0.8904],\n",
+      "        [ 0.9257],\n",
+      "        [ 0.9215],\n",
+      "        [ 0.8892],\n",
+      "        [ 0.8094],\n",
+      "        [ 0.6892],\n",
+      "        [ 0.5396],\n",
+      "        [ 0.3761],\n",
+      "        [ 0.2028],\n",
+      "        [ 0.0313],\n",
+      "        [-0.1376],\n",
+      "        [-0.2999]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "data_module.setup(stage=\"test\")\n",
+    "test_dataloader = data_module.test_dataloader()\n",
+    "\n",
+    "preds_mae = model_mae.predict(\n",
+    "    test_dataloader,\n",
+    "    mode=\"prediction\",\n",
+    "    trainer_kwargs={\"accelerator\": \"cpu\", \"enable_progress_bar\": False},\n",
+    ")\n",
+    "print(\"Point prediction shape:\", preds_mae[\"prediction\"].shape)\n",
+    "print(\"Point prediction:\", preds_mae[\"prediction\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a17e8f8b",
+   "metadata": {},
+   "source": [
+    "Same for the `QuantileLoss` model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "aaa6dc00",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (mps), used: False\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "TPU available: False, using: 0 TPU cores\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Quantile prediction shape (batch, pred_len, n_quantiles): torch.Size([255, 1, 3])\n",
+      "Quantile prediction: tensor([[[ 9.9807e-02,  2.7499e-01,  4.4650e-01]],\n",
+      "\n",
+      "        [[ 2.0286e-01,  4.5577e-01,  5.5673e-01]],\n",
+      "\n",
+      "        [[ 3.7001e-01,  5.8030e-01,  6.6626e-01]],\n",
+      "\n",
+      "        [[ 4.9127e-01,  7.1805e-01,  7.8477e-01]],\n",
+      "\n",
+      "        [[ 5.7864e-01,  7.9167e-01,  8.6310e-01]],\n",
+      "\n",
+      "        [[ 6.0197e-01,  8.7421e-01,  9.0665e-01]],\n",
+      "\n",
+      "        [[ 6.4056e-01,  9.0229e-01,  9.1842e-01]],\n",
+      "\n",
+      "        [[ 6.0831e-01,  8.7775e-01,  9.1337e-01]],\n",
+      "\n",
+      "        [[ 5.7948e-01,  8.5235e-01,  8.3199e-01]],\n",
+      "\n",
+      "        [[ 4.7244e-01,  7.5956e-01,  7.3842e-01]],\n",
+      "\n",
+      "        [[ 3.8413e-01,  6.6504e-01,  6.0270e-01]],\n",
+      "\n",
+      "        [[ 2.2860e-01,  5.3190e-01,  4.7483e-01]],\n",
+      "\n",
+      "        [[ 8.0860e-02,  3.8528e-01,  3.1630e-01]],\n",
+      "\n",
+      "        [[-6.6180e-02,  2.2654e-01,  1.6566e-01]],\n",
+      "\n",
+      "        [[-2.2870e-01,  3.9990e-02,  1.0635e-02]],\n",
+      "\n",
+      "        [[-4.5350e-01, -1.6615e-01, -1.6245e-01]],\n",
+      "\n",
+      "        [[-5.6181e-01, -3.0119e-01, -2.7940e-01]],\n",
+      "\n",
+      "        [[ 1.1397e-01,  3.1443e-01,  4.5570e-01]],\n",
+      "\n",
+      "        [[ 2.3875e-01,  4.5722e-01,  5.6592e-01]],\n",
+      "\n",
+      "        [[ 3.9594e-01,  6.1202e-01,  6.8070e-01]],\n",
+      "\n",
+      "        [[ 5.0909e-01,  7.2343e-01,  7.9554e-01]],\n",
+      "\n",
+      "        [[ 5.6820e-01,  8.2549e-01,  8.8526e-01]],\n",
+      "\n",
+      "        [[ 6.0274e-01,  8.7167e-01,  9.0742e-01]],\n",
+      "\n",
+      "        [[ 6.3060e-01,  9.1073e-01,  9.2737e-01]],\n",
+      "\n",
+      "        [[ 6.2148e-01,  8.6380e-01,  9.0087e-01]],\n",
+      "\n",
+      "        [[ 5.6531e-01,  8.4161e-01,  8.2112e-01]],\n",
+      "\n",
+      "        [[ 4.7440e-01,  7.6080e-01,  7.2122e-01]],\n",
+      "\n",
+      "        [[ 3.5833e-01,  6.6437e-01,  6.0321e-01]],\n",
+      "\n",
+      "        [[ 2.2504e-01,  5.3037e-01,  4.6037e-01]],\n",
+      "\n",
+      "        [[ 5.6557e-02,  3.6581e-01,  3.1029e-01]],\n",
+      "\n",
+      "        [[-6.2241e-02,  2.1161e-01,  1.4741e-01]],\n",
+      "\n",
+      "        [[-2.6160e-01,  9.4035e-04, -1.0257e-02]],\n",
+      "\n",
+      "        [[-4.4650e-01, -1.5157e-01, -1.6410e-01]],\n",
+      "\n",
+      "        [[-5.8261e-01, -3.1896e-01, -2.8520e-01]],\n",
+      "\n",
+      "        [[ 6.8224e-02,  2.6715e-01,  4.2531e-01]],\n",
+      "\n",
+      "        [[ 2.3004e-01,  4.4067e-01,  5.5171e-01]],\n",
+      "\n",
+      "        [[ 3.6883e-01,  5.7923e-01,  6.7432e-01]],\n",
+      "\n",
+      "        [[ 4.9196e-01,  7.2738e-01,  7.8510e-01]],\n",
+      "\n",
+      "        [[ 5.8808e-01,  8.1275e-01,  8.7540e-01]],\n",
+      "\n",
+      "        [[ 6.5047e-01,  8.7580e-01,  9.2384e-01]],\n",
+      "\n",
+      "        [[ 6.5986e-01,  9.3328e-01,  9.6060e-01]],\n",
+      "\n",
+      "        [[ 6.7055e-01,  9.2779e-01,  9.4879e-01]],\n",
+      "\n",
+      "        [[ 6.1443e-01,  8.8840e-01,  8.7499e-01]],\n",
+      "\n",
+      "        [[ 5.2145e-01,  8.1166e-01,  7.7975e-01]],\n",
+      "\n",
+      "        [[ 4.1959e-01,  7.0152e-01,  6.5948e-01]],\n",
+      "\n",
+      "        [[ 2.7865e-01,  5.7722e-01,  5.0482e-01]],\n",
+      "\n",
+      "        [[ 1.1372e-01,  3.9628e-01,  3.4282e-01]],\n",
+      "\n",
+      "        [[-5.8308e-02,  2.4688e-01,  1.7393e-01]],\n",
+      "\n",
+      "        [[-2.3804e-01,  5.6334e-02,  1.1986e-02]],\n",
+      "\n",
+      "        [[-4.3431e-01, -1.1985e-01, -1.5714e-01]],\n",
+      "\n",
+      "        [[-5.8965e-01, -3.2342e-01, -2.9989e-01]],\n",
+      "\n",
+      "        [[ 1.0963e-01,  3.1168e-01,  4.6216e-01]],\n",
+      "\n",
+      "        [[ 2.7842e-01,  4.8819e-01,  5.7088e-01]],\n",
+      "\n",
+      "        [[ 3.8399e-01,  6.1198e-01,  6.9901e-01]],\n",
+      "\n",
+      "        [[ 4.9422e-01,  7.1578e-01,  7.7966e-01]],\n",
+      "\n",
+      "        [[ 5.9026e-01,  8.3217e-01,  8.7925e-01]],\n",
+      "\n",
+      "        [[ 6.2625e-01,  8.6161e-01,  9.1631e-01]],\n",
+      "\n",
+      "        [[ 6.3366e-01,  8.7784e-01,  9.1430e-01]],\n",
+      "\n",
+      "        [[ 5.9915e-01,  8.8720e-01,  8.7842e-01]],\n",
+      "\n",
+      "        [[ 5.3163e-01,  8.0650e-01,  7.9515e-01]],\n",
+      "\n",
+      "        [[ 4.6380e-01,  7.5668e-01,  7.0034e-01]],\n",
+      "\n",
+      "        [[ 3.3820e-01,  6.1961e-01,  5.8047e-01]],\n",
+      "\n",
+      "        [[ 1.9239e-01,  4.9909e-01,  4.2081e-01]],\n",
+      "\n",
+      "        [[ 4.3815e-02,  3.3371e-01,  2.7643e-01]],\n",
+      "\n",
+      "        [[-1.5598e-01,  1.4883e-01,  1.0373e-01]],\n",
+      "\n",
+      "        [[-2.9957e-01, -1.2554e-02, -4.8678e-02]],\n",
+      "\n",
+      "        [[-4.4921e-01, -1.8426e-01, -1.7194e-01]],\n",
+      "\n",
+      "        [[-5.7228e-01, -3.4601e-01, -2.8218e-01]],\n",
+      "\n",
+      "        [[ 8.3240e-02,  2.8180e-01,  4.4251e-01]],\n",
+      "\n",
+      "        [[ 2.2670e-01,  4.5275e-01,  5.4464e-01]],\n",
+      "\n",
+      "        [[ 3.5842e-01,  5.7559e-01,  6.6990e-01]],\n",
+      "\n",
+      "        [[ 4.9171e-01,  7.0034e-01,  7.7952e-01]],\n",
+      "\n",
+      "        [[ 5.4633e-01,  7.8802e-01,  8.5538e-01]],\n",
+      "\n",
+      "        [[ 6.0242e-01,  8.6036e-01,  8.9241e-01]],\n",
+      "\n",
+      "        [[ 6.2993e-01,  8.9038e-01,  9.1905e-01]],\n",
+      "\n",
+      "        [[ 6.1258e-01,  8.7093e-01,  8.8347e-01]],\n",
+      "\n",
+      "        [[ 5.5138e-01,  8.2841e-01,  8.2372e-01]],\n",
+      "\n",
+      "        [[ 4.7361e-01,  7.6866e-01,  7.3051e-01]],\n",
+      "\n",
+      "        [[ 3.7874e-01,  6.5683e-01,  6.1570e-01]],\n",
+      "\n",
+      "        [[ 2.2295e-01,  5.2539e-01,  4.5199e-01]],\n",
+      "\n",
+      "        [[ 9.5070e-02,  3.7595e-01,  3.1711e-01]],\n",
+      "\n",
+      "        [[-8.5887e-02,  2.0173e-01,  1.4964e-01]],\n",
+      "\n",
+      "        [[-2.5123e-01,  2.3295e-02, -9.8590e-03]],\n",
+      "\n",
+      "        [[-4.3143e-01, -1.4560e-01, -1.4391e-01]],\n",
+      "\n",
+      "        [[-5.5544e-01, -3.0226e-01, -2.6597e-01]],\n",
+      "\n",
+      "        [[ 5.9465e-02,  2.5712e-01,  4.1841e-01]],\n",
+      "\n",
+      "        [[ 2.2534e-01,  4.2906e-01,  5.3150e-01]],\n",
+      "\n",
+      "        [[ 3.3970e-01,  5.6810e-01,  6.6276e-01]],\n",
+      "\n",
+      "        [[ 4.6560e-01,  6.7612e-01,  7.5248e-01]],\n",
+      "\n",
+      "        [[ 5.5287e-01,  7.9469e-01,  8.5845e-01]],\n",
+      "\n",
+      "        [[ 6.1629e-01,  8.6074e-01,  9.1417e-01]],\n",
+      "\n",
+      "        [[ 6.2908e-01,  8.9467e-01,  9.3152e-01]],\n",
+      "\n",
+      "        [[ 6.3811e-01,  9.0170e-01,  8.9450e-01]],\n",
+      "\n",
+      "        [[ 5.9622e-01,  8.4047e-01,  8.6211e-01]],\n",
+      "\n",
+      "        [[ 4.9510e-01,  7.9654e-01,  7.4223e-01]],\n",
+      "\n",
+      "        [[ 3.9885e-01,  6.7851e-01,  6.3447e-01]],\n",
+      "\n",
+      "        [[ 2.4988e-01,  5.6526e-01,  4.9061e-01]],\n",
+      "\n",
+      "        [[ 1.2074e-01,  3.7796e-01,  3.3539e-01]],\n",
+      "\n",
+      "        [[-7.4694e-02,  2.3034e-01,  1.7265e-01]],\n",
+      "\n",
+      "        [[-2.1436e-01,  6.2538e-02,  2.4097e-02]],\n",
+      "\n",
+      "        [[-3.9239e-01, -1.0717e-01, -1.1646e-01]],\n",
+      "\n",
+      "        [[-5.5500e-01, -2.7147e-01, -2.5749e-01]],\n",
+      "\n",
+      "        [[ 8.5695e-02,  2.6148e-01,  4.2282e-01]],\n",
+      "\n",
+      "        [[ 2.3940e-01,  4.5270e-01,  5.5483e-01]],\n",
+      "\n",
+      "        [[ 3.6641e-01,  5.7146e-01,  6.6453e-01]],\n",
+      "\n",
+      "        [[ 4.7790e-01,  7.2527e-01,  7.7900e-01]],\n",
+      "\n",
+      "        [[ 5.6712e-01,  8.0482e-01,  8.6703e-01]],\n",
+      "\n",
+      "        [[ 6.0549e-01,  8.8421e-01,  9.2169e-01]],\n",
+      "\n",
+      "        [[ 6.5790e-01,  8.9285e-01,  9.2020e-01]],\n",
+      "\n",
+      "        [[ 6.1083e-01,  8.9196e-01,  9.2065e-01]],\n",
+      "\n",
+      "        [[ 5.6070e-01,  8.4824e-01,  8.2934e-01]],\n",
+      "\n",
+      "        [[ 4.7873e-01,  7.7305e-01,  7.3288e-01]],\n",
+      "\n",
+      "        [[ 3.7758e-01,  6.8113e-01,  6.1767e-01]],\n",
+      "\n",
+      "        [[ 2.4868e-01,  5.4074e-01,  4.9440e-01]],\n",
+      "\n",
+      "        [[ 9.3219e-02,  3.8444e-01,  3.1803e-01]],\n",
+      "\n",
+      "        [[-8.1127e-02,  2.0306e-01,  1.6303e-01]],\n",
+      "\n",
+      "        [[-2.2893e-01,  4.1839e-02,  7.8391e-03]],\n",
+      "\n",
+      "        [[-4.2516e-01, -1.3712e-01, -1.4634e-01]],\n",
+      "\n",
+      "        [[-5.4870e-01, -2.8975e-01, -2.7055e-01]],\n",
+      "\n",
+      "        [[ 8.9124e-02,  2.7928e-01,  4.3621e-01]],\n",
+      "\n",
+      "        [[ 2.3139e-01,  4.4494e-01,  5.5406e-01]],\n",
       "\n",
-      "        [[ 0.0680,  0.0936,  0.2504]],\n",
+      "        [[ 3.5274e-01,  5.8739e-01,  6.5556e-01]],\n",
       "\n",
-      "        [[ 0.2310,  0.2605,  0.4298]],\n",
+      "        [[ 4.4339e-01,  7.2095e-01,  7.6852e-01]],\n",
       "\n",
-      "        [[ 0.3604,  0.3968,  0.5679]],\n",
+      "        [[ 5.6372e-01,  8.1255e-01,  8.6896e-01]],\n",
       "\n",
-      "        [[ 0.4935,  0.5408,  0.7165]],\n",
+      "        [[ 6.1433e-01,  8.7696e-01,  9.2724e-01]],\n",
       "\n",
-      "        [[ 0.6274,  0.6697,  0.8745]],\n",
+      "        [[ 6.5213e-01,  8.9790e-01,  9.4222e-01]],\n",
       "\n",
-      "        [[ 0.7192,  0.7940,  0.9812]],\n",
+      "        [[ 6.4048e-01,  9.0298e-01,  9.1548e-01]],\n",
       "\n",
-      "        [[ 0.7555,  0.8650,  1.0313]],\n",
+      "        [[ 5.7732e-01,  8.5814e-01,  8.4854e-01]],\n",
       "\n",
-      "        [[ 0.7602,  0.8706,  1.0427]],\n",
+      "        [[ 4.9995e-01,  7.9472e-01,  7.5559e-01]],\n",
       "\n",
-      "        [[ 0.7532,  0.8524,  1.0308]],\n",
+      "        [[ 4.1071e-01,  7.1170e-01,  6.3694e-01]],\n",
       "\n",
-      "        [[ 0.7003,  0.7784,  0.9995]],\n",
+      "        [[ 2.7086e-01,  5.4928e-01,  4.9757e-01]],\n",
       "\n",
-      "        [[ 0.5987,  0.6807,  0.9390]],\n",
+      "        [[ 1.0377e-01,  3.9813e-01,  3.2809e-01]],\n",
       "\n",
-      "        [[ 0.4757,  0.5814,  0.7966]],\n",
+      "        [[-7.8528e-02,  2.3026e-01,  1.6900e-01]],\n",
       "\n",
-      "        [[ 0.3432,  0.4587,  0.6614]],\n",
+      "        [[-2.3633e-01,  5.8970e-02,  1.2067e-02]],\n",
       "\n",
-      "        [[ 0.1659,  0.2931,  0.5039]],\n",
+      "        [[-3.5787e-01, -9.6255e-02, -1.1465e-01]],\n",
       "\n",
-      "        [[-0.0338,  0.0983,  0.3208]],\n",
+      "        [[-5.3583e-01, -2.8372e-01, -2.4692e-01]],\n",
       "\n",
-      "        [[-0.1989, -0.0829,  0.1821]],\n",
+      "        [[ 6.4040e-02,  2.6271e-01,  4.1785e-01]],\n",
       "\n",
-      "        [[-0.3732, -0.2402,  0.0121]],\n",
+      "        [[ 2.1890e-01,  4.0835e-01,  5.3605e-01]],\n",
       "\n",
-      "        [[-0.5151, -0.3600, -0.1606]],\n",
+      "        [[ 3.3648e-01,  5.6810e-01,  6.6188e-01]],\n",
       "\n",
-      "        [[-0.0789, -0.0406,  0.0908]],\n",
+      "        [[ 4.5166e-01,  6.7471e-01,  7.6441e-01]],\n",
       "\n",
-      "        [[ 0.0495,  0.0830,  0.2585]],\n",
+      "        [[ 5.3114e-01,  7.8730e-01,  8.4444e-01]],\n",
       "\n",
-      "        [[ 0.2185,  0.2520,  0.4223]],\n",
+      "        [[ 6.2036e-01,  8.5816e-01,  8.9598e-01]],\n",
       "\n",
-      "        [[ 0.3870,  0.4209,  0.5818]],\n",
+      "        [[ 6.2312e-01,  8.8417e-01,  9.2671e-01]],\n",
       "\n",
-      "        [[ 0.5243,  0.5766,  0.7636]],\n",
+      "        [[ 6.3499e-01,  8.9551e-01,  8.9948e-01]],\n",
       "\n",
-      "        [[ 0.6293,  0.6854,  0.8715]],\n",
+      "        [[ 5.8168e-01,  8.5119e-01,  8.5915e-01]],\n",
       "\n",
-      "        [[ 0.7055,  0.7854,  0.9698]],\n",
+      "        [[ 4.9857e-01,  7.8948e-01,  7.5264e-01]],\n",
       "\n",
-      "        [[ 0.7722,  0.8390,  1.0474]],\n",
+      "        [[ 3.9344e-01,  6.6870e-01,  6.3241e-01]],\n",
       "\n",
-      "        [[ 0.8323,  0.9074,  1.0969]],\n",
+      "        [[ 2.5457e-01,  5.4290e-01,  4.7884e-01]],\n",
       "\n",
-      "        [[ 0.8132,  0.8968,  1.1051]],\n",
+      "        [[ 1.0606e-01,  3.8402e-01,  3.2660e-01]],\n",
       "\n",
-      "        [[ 0.6892,  0.8067,  1.0172]],\n",
+      "        [[-5.6515e-02,  2.2615e-01,  1.6695e-01]],\n",
       "\n",
-      "        [[ 0.5896,  0.7130,  0.9167]],\n",
+      "        [[-2.4996e-01,  4.9497e-02,  4.9640e-03]],\n",
       "\n",
-      "        [[ 0.4989,  0.5976,  0.8067]]])\n"
+      "        [[-4.2720e-01, -1.1386e-01, -1.3620e-01]],\n",
+      "\n",
+      "        [[-5.6947e-01, -2.9031e-01, -2.6498e-01]],\n",
+      "\n",
+      "        [[ 9.5746e-02,  3.0031e-01,  4.5950e-01]],\n",
+      "\n",
+      "        [[ 2.5609e-01,  4.8182e-01,  5.6761e-01]],\n",
+      "\n",
+      "        [[ 3.8856e-01,  6.1652e-01,  7.1072e-01]],\n",
+      "\n",
+      "        [[ 4.9245e-01,  7.1326e-01,  7.8025e-01]],\n",
+      "\n",
+      "        [[ 5.6165e-01,  8.0840e-01,  8.5138e-01]],\n",
+      "\n",
+      "        [[ 6.3295e-01,  8.7799e-01,  9.2276e-01]],\n",
+      "\n",
+      "        [[ 6.2531e-01,  8.9312e-01,  9.3346e-01]],\n",
+      "\n",
+      "        [[ 6.0895e-01,  8.8396e-01,  8.8771e-01]],\n",
+      "\n",
+      "        [[ 5.5706e-01,  8.3674e-01,  8.2919e-01]],\n",
+      "\n",
+      "        [[ 4.7386e-01,  7.5375e-01,  7.2319e-01]],\n",
+      "\n",
+      "        [[ 3.5052e-01,  6.3204e-01,  5.8431e-01]],\n",
+      "\n",
+      "        [[ 2.0327e-01,  5.3813e-01,  4.4271e-01]],\n",
+      "\n",
+      "        [[ 8.6660e-02,  3.5091e-01,  3.0531e-01]],\n",
+      "\n",
+      "        [[-1.0494e-01,  1.9784e-01,  1.3496e-01]],\n",
+      "\n",
+      "        [[-2.8547e-01,  6.0597e-03, -2.6737e-02]],\n",
+      "\n",
+      "        [[-4.4731e-01, -1.6901e-01, -1.7081e-01]],\n",
+      "\n",
+      "        [[-5.6197e-01, -3.3158e-01, -2.8447e-01]],\n",
+      "\n",
+      "        [[ 1.0085e-01,  2.8361e-01,  4.3741e-01]],\n",
+      "\n",
+      "        [[ 2.5792e-01,  4.5220e-01,  5.5045e-01]],\n",
+      "\n",
+      "        [[ 3.4085e-01,  5.7949e-01,  6.6339e-01]],\n",
+      "\n",
+      "        [[ 5.0523e-01,  7.1858e-01,  7.6561e-01]],\n",
+      "\n",
+      "        [[ 5.8831e-01,  8.1191e-01,  8.9661e-01]],\n",
+      "\n",
+      "        [[ 6.2844e-01,  8.6103e-01,  9.0313e-01]],\n",
+      "\n",
+      "        [[ 6.4458e-01,  9.0609e-01,  9.4090e-01]],\n",
+      "\n",
+      "        [[ 6.3490e-01,  8.9844e-01,  9.0340e-01]],\n",
+      "\n",
+      "        [[ 5.5762e-01,  8.3458e-01,  8.2804e-01]],\n",
+      "\n",
+      "        [[ 4.8105e-01,  7.7549e-01,  7.2353e-01]],\n",
+      "\n",
+      "        [[ 3.7026e-01,  6.7048e-01,  6.1581e-01]],\n",
+      "\n",
+      "        [[ 2.2807e-01,  5.1181e-01,  4.4185e-01]],\n",
+      "\n",
+      "        [[ 3.7930e-02,  3.6447e-01,  2.9856e-01]],\n",
+      "\n",
+      "        [[-1.1246e-01,  2.0081e-01,  1.3358e-01]],\n",
+      "\n",
+      "        [[-2.6385e-01,  6.7742e-03, -1.1569e-02]],\n",
+      "\n",
+      "        [[-4.5742e-01, -1.8700e-01, -1.8088e-01]],\n",
+      "\n",
+      "        [[-5.9037e-01, -3.5330e-01, -2.8659e-01]],\n",
+      "\n",
+      "        [[ 6.5153e-02,  2.3592e-01,  4.0362e-01]],\n",
+      "\n",
+      "        [[ 1.7578e-01,  3.8005e-01,  5.1560e-01]],\n",
+      "\n",
+      "        [[ 2.9644e-01,  5.2635e-01,  6.1373e-01]],\n",
+      "\n",
+      "        [[ 4.3944e-01,  6.6463e-01,  7.4343e-01]],\n",
+      "\n",
+      "        [[ 5.0876e-01,  7.3946e-01,  8.0405e-01]],\n",
+      "\n",
+      "        [[ 5.9904e-01,  8.2615e-01,  8.8129e-01]],\n",
+      "\n",
+      "        [[ 6.1663e-01,  8.7560e-01,  9.1641e-01]],\n",
+      "\n",
+      "        [[ 6.2345e-01,  8.7698e-01,  9.0942e-01]],\n",
+      "\n",
+      "        [[ 5.5821e-01,  8.4739e-01,  8.3576e-01]],\n",
+      "\n",
+      "        [[ 5.1653e-01,  7.9141e-01,  7.6143e-01]],\n",
+      "\n",
+      "        [[ 4.1469e-01,  6.7007e-01,  6.4260e-01]],\n",
+      "\n",
+      "        [[ 2.7160e-01,  5.6566e-01,  4.9475e-01]],\n",
+      "\n",
+      "        [[ 1.1492e-01,  4.1242e-01,  3.5776e-01]],\n",
+      "\n",
+      "        [[-3.8015e-02,  2.5949e-01,  2.0511e-01]],\n",
+      "\n",
+      "        [[-2.0979e-01,  8.4884e-02,  4.2242e-02]],\n",
+      "\n",
+      "        [[-3.6997e-01, -8.6048e-02, -1.0172e-01]],\n",
+      "\n",
+      "        [[-5.1088e-01, -2.4791e-01, -2.2313e-01]],\n",
+      "\n",
+      "        [[ 9.8262e-02,  3.1037e-01,  4.5498e-01]],\n",
+      "\n",
+      "        [[ 2.5056e-01,  4.9253e-01,  6.0234e-01]],\n",
+      "\n",
+      "        [[ 3.7600e-01,  6.4038e-01,  6.9735e-01]],\n",
+      "\n",
+      "        [[ 5.0658e-01,  7.4425e-01,  7.9832e-01]],\n",
+      "\n",
+      "        [[ 5.9397e-01,  8.4450e-01,  8.9360e-01]],\n",
+      "\n",
+      "        [[ 6.3720e-01,  8.8409e-01,  9.2933e-01]],\n",
+      "\n",
+      "        [[ 6.3565e-01,  9.1066e-01,  9.4339e-01]],\n",
+      "\n",
+      "        [[ 6.3572e-01,  9.1659e-01,  9.2350e-01]],\n",
+      "\n",
+      "        [[ 5.8792e-01,  8.4999e-01,  8.4288e-01]],\n",
+      "\n",
+      "        [[ 4.8319e-01,  7.7887e-01,  7.3816e-01]],\n",
+      "\n",
+      "        [[ 3.8246e-01,  6.7503e-01,  6.1371e-01]],\n",
+      "\n",
+      "        [[ 2.3944e-01,  5.3841e-01,  4.7206e-01]],\n",
+      "\n",
+      "        [[ 7.9835e-02,  3.8465e-01,  3.1268e-01]],\n",
+      "\n",
+      "        [[-7.1891e-02,  2.2273e-01,  1.6437e-01]],\n",
+      "\n",
+      "        [[-2.5790e-01,  2.9014e-02, -1.1495e-02]],\n",
+      "\n",
+      "        [[-4.2010e-01, -1.5221e-01, -1.6836e-01]],\n",
+      "\n",
+      "        [[-5.5266e-01, -2.9643e-01, -2.7069e-01]],\n",
+      "\n",
+      "        [[ 6.5341e-02,  2.7252e-01,  4.2122e-01]],\n",
+      "\n",
+      "        [[ 2.5403e-01,  4.7456e-01,  5.6344e-01]],\n",
+      "\n",
+      "        [[ 3.9813e-01,  6.0208e-01,  6.9627e-01]],\n",
+      "\n",
+      "        [[ 4.8692e-01,  7.1395e-01,  7.9074e-01]],\n",
+      "\n",
+      "        [[ 5.5250e-01,  8.2135e-01,  8.6205e-01]],\n",
+      "\n",
+      "        [[ 6.1700e-01,  8.7382e-01,  9.1996e-01]],\n",
+      "\n",
+      "        [[ 6.3738e-01,  8.9887e-01,  9.3028e-01]],\n",
+      "\n",
+      "        [[ 6.2956e-01,  8.8642e-01,  8.9588e-01]],\n",
+      "\n",
+      "        [[ 5.8343e-01,  8.5662e-01,  8.4109e-01]],\n",
+      "\n",
+      "        [[ 5.0555e-01,  7.7396e-01,  7.5080e-01]],\n",
+      "\n",
+      "        [[ 3.5991e-01,  6.6985e-01,  6.0834e-01]],\n",
+      "\n",
+      "        [[ 2.2618e-01,  5.4520e-01,  4.6742e-01]],\n",
+      "\n",
+      "        [[ 8.8920e-02,  3.5737e-01,  3.0228e-01]],\n",
+      "\n",
+      "        [[-9.1688e-02,  1.7764e-01,  1.2848e-01]],\n",
+      "\n",
+      "        [[-2.9603e-01,  1.1382e-02, -2.4436e-02]],\n",
+      "\n",
+      "        [[-4.3416e-01, -1.2885e-01, -1.4868e-01]],\n",
+      "\n",
+      "        [[-5.6478e-01, -2.8707e-01, -2.5098e-01]],\n",
+      "\n",
+      "        [[ 8.6806e-02,  2.9878e-01,  4.3477e-01]],\n",
+      "\n",
+      "        [[ 2.4188e-01,  4.4121e-01,  5.5636e-01]],\n",
+      "\n",
+      "        [[ 3.6820e-01,  6.0269e-01,  6.6721e-01]],\n",
+      "\n",
+      "        [[ 4.6699e-01,  6.7692e-01,  7.7065e-01]],\n",
+      "\n",
+      "        [[ 5.4361e-01,  8.1212e-01,  8.5542e-01]],\n",
+      "\n",
+      "        [[ 5.9484e-01,  8.5596e-01,  9.0012e-01]],\n",
+      "\n",
+      "        [[ 6.3045e-01,  8.8823e-01,  9.2780e-01]],\n",
+      "\n",
+      "        [[ 6.2737e-01,  8.7979e-01,  8.9424e-01]],\n",
+      "\n",
+      "        [[ 5.6011e-01,  8.4219e-01,  8.4310e-01]],\n",
+      "\n",
+      "        [[ 4.9571e-01,  7.8150e-01,  7.3383e-01]],\n",
+      "\n",
+      "        [[ 3.7415e-01,  6.6597e-01,  6.1190e-01]],\n",
+      "\n",
+      "        [[ 2.3752e-01,  5.3129e-01,  4.6802e-01]],\n",
+      "\n",
+      "        [[ 7.8355e-02,  3.8520e-01,  3.1256e-01]],\n",
+      "\n",
+      "        [[-5.8871e-02,  2.0440e-01,  1.6026e-01]],\n",
+      "\n",
+      "        [[-2.5583e-01,  2.8542e-02,  3.4320e-03]],\n",
+      "\n",
+      "        [[-3.9615e-01, -1.2049e-01, -1.3581e-01]],\n",
+      "\n",
+      "        [[-5.4790e-01, -2.8266e-01, -2.4942e-01]]])\n"
      ]
     }
    ],
    "source": [
-    "with torch.no_grad():\n",
-    "    test_batch = next(iter(data_module.test_dataloader()))\n",
-    "    x_test, y_test = test_batch\n",
-    "    y_pred = model2(x_test)\n",
-    "\n",
-    "    print(\"Prediction:\", y_pred[\"prediction\"])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "id": "a4e6e4b1",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "torch.Size([32, 1, 3])"
-      ]
-     },
-     "execution_count": 19,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "y_pred[\"prediction\"].shape"
+    "preds_q = model_q.predict(\n",
+    "    test_dataloader,\n",
+    "    mode=\"quantiles\",\n",
+    "    trainer_kwargs={\"accelerator\": \"cpu\", \"enable_progress_bar\": False},\n",
+    ")\n",
+    "print(\n",
+    "    \"Quantile prediction shape (batch, pred_len, n_quantiles):\",\n",
+    "    preds_q[\"prediction\"].shape,\n",
+    ")\n",
+    "print(\"Quantile prediction:\", preds_q[\"prediction\"])"
    ]
   }
  ],
@@ -1269,7 +1995,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.3"
+   "version": "3.14.6"
   }
  },
  "nbformat": 4,