diff --git a/README.md b/README.md
index f78cf0c..f711193 100644
--- a/README.md
+++ b/README.md
@@ -60,12 +60,13 @@ loader = DataLoader(ds, batch_size=8, collate_fn=neural_collate)
 # 2) Load a pretrained model from the Hugging Face Hub.
 model = StateNet.from_pretrained("urancon/deepSTRF-statenet-gru-ns1").eval()
 
-# 3) Score it.
-stims, responses, _, _ = next(iter(loader))
-pred    = model(stims)                                      # (B, N, R=1, T)
-psth    = responses.nanmean(dim=2, keepdim=True)
-cc      = corrcoef(pred, psth, reduction='mean')
-cc_norm = normalized_corrcoef(pred, responses, method='schoppe', reduction='mean')
+# 3) Score it. Each batch is a dict: 'stims', 'responses', 'valid_mask', 'stim_meta'.
+batch     = next(iter(loader))
+responses = batch['responses']
+pred      = model(batch['stims'])                           # (B, N, R=1, T)
+psth      = responses.nanmean(dim=2, keepdim=True)
+cc        = corrcoef(pred, psth, reduction='mean')
+cc_norm   = normalized_corrcoef(pred, responses, method='schoppe', reduction='mean')
 print(f"CCraw = {cc:.3f}   CCnorm = {cc_norm:.3f}")
 ```
 
diff --git a/deepSTRF/datasets/neural_dataset.py b/deepSTRF/datasets/neural_dataset.py
index d0f5e85..a8fafed 100644
--- a/deepSTRF/datasets/neural_dataset.py
+++ b/deepSTRF/datasets/neural_dataset.py
@@ -260,9 +260,14 @@ def __getitem__(self, idx):
 
         Returns
         -------
-        stim, responses, mask, stim_meta
-            Tuple for a scalar index, or 4 lists for slice / list indexing.
-            ``responses`` and ``mask`` are restricted to the selected neurons.
+        dict
+            A dict with keys ``'stims'``, ``'responses'``, ``'valid_mask'``,
+            ``'stim_meta'``. For a scalar index each value is a single item;
+            for slice / list indexing each value is a list. ``'responses'``
+            and ``'valid_mask'`` are restricted to the selected neurons. The
+            dict container (rather than a positional tuple) lets datasets add
+            extra per-trial keys later — e.g. ``'behav'`` for behavioural
+            covariates — without changing the unpacking contract.
         """
         iter_idx = self._iter_idx  # snapshot once; O(S * |I|) per call
 
@@ -292,8 +297,10 @@ def __getitem__(self, idx):
         masks = [all_masks[i][selected] for i in indices]
 
         if single:
-            return stims[0], resps[0], masks[0], metas[0]
-        return stims, resps, masks, metas
+            return {'stims': stims[0], 'responses': resps[0],
+                    'valid_mask': masks[0], 'stim_meta': metas[0]}
+        return {'stims': stims, 'responses': resps,
+                'valid_mask': masks, 'stim_meta': metas}
 
     def __repr__(self):
         return (f"{self.__class__.__name__}(N_neurons={self.get_N()}, "
diff --git a/deepSTRF/training/fitter.py b/deepSTRF/training/fitter.py
index f9d4cee..e0ccea9 100644
--- a/deepSTRF/training/fitter.py
+++ b/deepSTRF/training/fitter.py
@@ -376,9 +376,8 @@ def _train_one_epoch(self) -> Dict[str, Any]:
         responses_list: List[torch.Tensor] = []
 
         for batch in self.train_loader:
-            stims, responses, _valid_mask, _stim_metas = batch
-            stims = stims.to(self.device)
-            responses = responses.to(self.device)
+            stims = batch['stims'].to(self.device)
+            responses = batch['responses'].to(self.device)
 
             self.optimizer.zero_grad()
             pred = self.model(stims)
@@ -417,9 +416,8 @@ def _evaluate(self, loader: DataLoader) -> Dict[str, Any]:
 
         with torch.no_grad():
             for batch in loader:
-                stims, responses, _valid_mask, _stim_metas = batch
-                stims = stims.to(self.device)
-                responses = responses.to(self.device)
+                stims = batch['stims'].to(self.device)
+                responses = batch['responses'].to(self.device)
 
                 pred = self.model(stims)
                 if hasattr(self.model, "detach"):
diff --git a/deepSTRF/utils/data.py b/deepSTRF/utils/data.py
index 6fee65b..560b924 100644
--- a/deepSTRF/utils/data.py
+++ b/deepSTRF/utils/data.py
@@ -109,43 +109,45 @@ def neural_collate(batch):
 
     Parameters
     ----------
-    batch : list of 4-tuples
-        Each tuple is ``(stim, per_neuron_responses, per_neuron_mask,
-        stim_meta)`` as yielded by ``NeuralDataset.__getitem__`` for a
-        single item:
+    batch : list of dict
+        Each dict is one item as yielded by ``NeuralDataset.__getitem__``,
+        with keys:
 
-        * ``stim`` — a stim tensor of shape ``(..., T_s)`` (modality-specific
+        * ``'stims'`` — a stim tensor of shape ``(..., T_s)`` (modality-specific
           leading dims, e.g. ``(1, F, T_s)`` for audio).
-        * ``per_neuron_responses`` — list of length ``N_selected``; each
-          element is a ``(R_{s,n}, T_s)`` spike-count tensor or a
-          ``(1, 1)`` NaN sentinel.
-        * ``per_neuron_mask`` — ``(N_selected,)`` bool tensor (currently
-          ignored; the fine-grained ``valid_mask`` returned by this
-          function subsumes it).
-        * ``stim_meta`` — per-stim metadata dict.
+        * ``'responses'`` — list of length ``N_selected``; each element is a
+          ``(R_{s,n}, T_s)`` spike-count tensor or a ``(1, 1)`` NaN sentinel.
+        * ``'valid_mask'`` — ``(N_selected,)`` per-neuron bool tensor (ignored
+          here; the fine-grained batch ``'valid_mask'`` below subsumes it).
+        * ``'stim_meta'`` — per-stim metadata dict.
+
+        Extra keys (e.g. ``'behav'``) are passed through untouched: any key
+        not handled explicitly is collected into a length-``B`` list.
 
     Returns
     -------
-    stims : torch.Tensor
-        ``(B, ..., T_stim_max)`` float tensor, zero-padded along the last
-        axis. Contains no NaN.
-    responses : torch.Tensor
-        ``(B, N_selected, R_max, T_resp_max)`` float tensor. NaN-padded
-        along both the repeat (``R``) and time (``T``) axes. Fully-NaN
-        slabs mark (stim, neuron) pairs with no recorded data. The
-        response-time axis is sized to ``T_resp_max`` independently of the
-        stim-time axis: in spectrogram mode the two are equal (one bin
-        per neural sample), but in waveform mode the stim axis runs at
-        ``audio_fs`` Hz while responses stay at the dataset's neural
-        ``dt_ms`` rate.
-    valid_mask : torch.Tensor
-        ``(B, N_selected, R_max, T_resp_max)`` bool tensor.
-        ``~responses.isnan()``, cached here so downstream loss code does
-        not have to recompute.
-    stim_metas : list
-        Length-``B`` list of the per-item stim_meta dicts.
+    dict
+        A dict with keys:
+
+        * ``'stims'`` — ``(B, ..., T_stim_max)`` float tensor, zero-padded
+          along the last axis. Contains no NaN.
+        * ``'responses'`` — ``(B, N_selected, R_max, T_resp_max)`` float
+          tensor. NaN-padded along both the repeat (``R``) and time (``T``)
+          axes. Fully-NaN slabs mark (stim, neuron) pairs with no recorded
+          data. The response-time axis is sized to ``T_resp_max``
+          independently of the stim-time axis: in spectrogram mode the two
+          are equal (one bin per neural sample), but in waveform mode the
+          stim axis runs at ``audio_fs`` Hz while responses stay at the
+          dataset's neural ``dt_ms`` rate.
+        * ``'valid_mask'`` — ``(B, N_selected, R_max, T_resp_max)`` bool
+          tensor, ``~responses.isnan()``, cached here so downstream loss code
+          does not have to recompute.
+        * ``'stim_meta'`` — length-``B`` list of the per-item stim_meta dicts.
+        * any extra per-item keys — length-``B`` lists, passed through.
     """
-    stims_list, resps_list, _masks_list, metas_list = zip(*batch)
+    stims_list = [item['stims'] for item in batch]
+    resps_list = [item['responses'] for item in batch]
+    metas_list = [item['stim_meta'] for item in batch]
     B = len(stims_list)
     N = len(resps_list[0])
 
@@ -171,7 +173,21 @@ def neural_collate(batch):
     # "for free" and does not need to scan again.
     valid_mask = ~responses.isnan()
 
-    return stims, responses, valid_mask, list(metas_list)
+    out = {
+        'stims': stims,
+        'responses': responses,
+        'valid_mask': valid_mask,
+        'stim_meta': list(metas_list),
+    }
+
+    # pass through any extra per-item keys (e.g. 'behav') as length-B lists,
+    # so datasets can add covariates without touching this collate.
+    handled = {'stims', 'responses', 'valid_mask', 'stim_meta'}
+    for key in batch[0]:
+        if key not in handled:
+            out[key] = [item[key] for item in batch]
+
+    return out
 
 
 def concat_neural_datasets(datasets: Sequence[NeuralDataset],
diff --git a/docs/_source/md/README_datasets.md b/docs/_source/md/README_datasets.md
index 4edb67f..11599d4 100644
--- a/docs/_source/md/README_datasets.md
+++ b/docs/_source/md/README_datasets.md
@@ -1,37 +1,71 @@
 # Conventions
 
-One of the main contributions of this repository is the presence of several off-the-shelf electrophysiology datasets, 
-compiled from various sources, preprocessed, and represented by convenient PyTorch classes.
-
-A major limitation hindering research in the field of auditory neural response fitting, is the great variability of 
-preprocessing methods and formats, due to signals of different nature (e.g. patch-clamp, extracellular recordings),
-different groups having different habits, different experimental setups, etc. This is why we tried to adopt a single 
-"format" for all datasets present in this repository.
-
-All datasets inherit from `torch.utils.data.dataset.Dataset` and can therefore benefit from other PyTorch utilities, 
-such as random splitting, shuffling, concatenating, augmentation, and so on.
-
-Because original data differ a lot between sources, the core and attributes of each dataset class might differ a bit, 
-but they still share a lot in common and can be used in very similar ways from the exterior.
-Namely, the constructor of dataset classes must include one argument to select neurons according to a criterion (mask, 
-index, a label, brain area). By default, all neurons are selected and teh constructor will grab the data for all of them. Similarly, the ability to choose stimuli of different nature (natural, pure tone, chord, etc.) 
-should be implemented as well. If several signals are available for each neuron, then an option should be made to choose 
-between them.
-
-In this framework, each dataset contain different sets of data for each neuron, notably:
-* input stimuli (spectrograms) that were presented to this neuron
-* the corresponding response, aligned in time
-* per-stimulus metrics, like measures of variability (e.g., *TTRC*, *CC_half*) pre-computed once during dataset instanciation,
-to save time during fitting
-
-As mentioned above, the set of units to work with is defined according to the user's criterion, in the constructor. 
-This number of neurons becomes the attribute `self.N_neurons`. 
-
-An important thing to note is that depending on the dataset, neurons may have different numbers of stimulus/response pairs,
-of different durations. Because of this, it can be necessary to fit computational models *unit by unit*, and then average 
-results across units. Therefore, all datasets come with a method `select_neuron(index)` that sets the attribute `self.I` 
-to the desired `index` (0 by default at instanciation). This attribute can also be a list of multiple desired index, if
-you want to train on a *population*, rather than a single isolated unit. You can then iterate over the stimulus/response 
-pairs of the currently selected neurons.
-
-Each call of the dataset or associated dataloader will return 4 things: `spectrogram, responses, ccmax, ttrc`.
+One of the main contributions of this library is a growing zoo of off-the-shelf
+electrophysiology (and EEG) datasets, compiled from various public sources,
+preprocessed, and exposed through a single PyTorch API.
+
+A major obstacle in sensory-response fitting is the sheer variability of
+preprocessing methods and formats — different signals (extracellular spikes,
+multi-unit activity, intracellular potential, scalp EEG), different labs, and
+different experimental setups. deepSTRF hides that behind one common contract so
+the same model and training loop work across datasets.
+
+## One base class, one shape
+
+Every dataset subclasses
+{class}`~deepSTRF.datasets.neural_dataset.NeuralDataset` (itself a
+`torch.utils.data.Dataset`), so it composes with the usual PyTorch utilities
+(splitting, shuffling, concatenation, …). Internally each dataset stores, per
+stimulus:
+
+* the **stimulus** presented (a spectrogram `(1, F, T)`, or a raw waveform
+  `(1, T_audio)` when `return_waveform=True`);
+* the **trial-resolved responses**, time-aligned to the stimulus;
+* per-stimulus and per-neuron **metadata** dicts (`stim_meta`, `nrn_meta`).
+
+Datasets are sparse and ragged — neurons may not all hear every stimulus, and
+stimuli vary in duration and repeat count. deepSTRF handles this with **NaN
+sentinels** for missing (stimulus, neuron) trials and zero-/NaN-padding at
+collate time. The full contract — the canonical `(B, N, R, T)` response shape,
+the NaN-sentinel rules, and the bidirectional neuron/stimulus selection — is
+documented in {doc}`data_paradigm`. Read that before touching any
+response-path code.
+
+## What a batch looks like
+
+Pair a dataset with `neural_collate` and a `DataLoader`. Each batch is a
+**dict** (not a positional tuple), so future per-trial variables can be added as
+new keys without breaking your unpacking:
+
+```python
+from torch.utils.data import DataLoader
+from deepSTRF.utils.data import neural_collate
+
+loader = DataLoader(ds, batch_size=8, collate_fn=neural_collate)
+
+for batch in loader:
+    stims      = batch['stims']        # (B, ..., T)        zero-padded, no NaN
+    responses  = batch['responses']    # (B, N, R, T)       NaN-padded
+    valid_mask = batch['valid_mask']   # (B, N, R, T) bool  ~responses.isnan()
+    metas      = batch['stim_meta']    # length-B list of per-stim dicts
+    ...
+```
+
+Indexing the dataset directly (`ds[i]`) returns the same keys for a single item.
+`CCmax` / `TTRC`-style normalisation is **not** precomputed at the dataloader
+boundary — it is derived on demand from `responses` by the metrics in
+{doc}`metrics_paradigm` (e.g. `normalized_corrcoef`).
+
+## Selecting neurons and stimuli
+
+The set of units (and stimuli) to work with is chosen through the selection API
+rather than a fixed constructor argument — by default all neurons are selected.
+The selection drives both `len(ds)` and iteration. See {doc}`data_paradigm` for
+the exact semantics; the common entry points are:
+
+```python
+ds.select_population([0, 1, 2])                          # by index (or a single int)
+ds.select_pop_by_nrn_attr("area", "Field_L")            # by metadata label
+ds.select_pop_by_nrn_predicate(lambda n: n.get("snr", 0) > 0.5)   # by threshold
+ds.select_stims_by_attr("type", "human_speech")         # restrict the stimulus set
+```
diff --git a/docs/_source/md/data_paradigm.md b/docs/_source/md/data_paradigm.md
index 85e6f8a..fb856b8 100644
--- a/docs/_source/md/data_paradigm.md
+++ b/docs/_source/md/data_paradigm.md
@@ -219,7 +219,8 @@ NaN-aware by default. The training loop therefore stays simple:
 ```python
 from deepSTRF.metrics import mse_loss, corrcoef, normalized_corrcoef
 
-for stims, responses, valid_mask, stim_metas in loader:
+for batch in loader:                                      # batch is a dict with keys
+    stims, responses = batch['stims'], batch['responses'] #   'stims','responses','valid_mask','stim_meta'
     pred    = model(stims)                                # (B, N, 1, T_max)
     gt_psth = responses.nanmean(dim=2, keepdim=True)      # (B, N, 1, T_max)
 
diff --git a/docs/_source/md/fitter.md b/docs/_source/md/fitter.md
index c351c28..5218ef3 100644
--- a/docs/_source/md/fitter.md
+++ b/docs/_source/md/fitter.md
@@ -21,9 +21,9 @@ canonical path is the three-line loop documented in
 `metrics_paradigm.md` §7:
 
 ```python
-for stims, responses, valid_mask, stim_metas in loader:
-    pred = model(stims)                                    # (B, N, 1, T)
-    loss = mse_loss(pred, responses)                       # auto-PSTH inside
+for batch in loader:                                       # batch is a dict
+    pred = model(batch['stims'])                           # (B, N, 1, T)
+    loss = mse_loss(pred, batch['responses'])              # auto-PSTH inside
     loss.backward(); optimizer.step(); optimizer.zero_grad()
 ```
 
@@ -553,11 +553,12 @@ The legacy module exposes `optimize_multiple_seeds`,
 | `set_random_seed`            | `deepSTRF.training.set_random_seed` (kept as deprecated re-export)  |
 
 Behavioral differences worth flagging:
-- The legacy code unpacks `(spectrogram, responses, ccmax, ttrc)`. The
-  new code unpacks `(stims, responses, valid_mask, stim_metas)` from
-  `neural_collate`. `ccmax` and `ttrc` are no longer dataloader-side
-  pre-computed tensors — they are computed on demand by
-  `normalized_corrcoef` from raw `responses`.
+- The legacy code unpacks a `(spectrogram, responses, ccmax, ttrc)`
+  tuple. `neural_collate` now yields a **dict** with keys `'stims'`,
+  `'responses'`, `'valid_mask'`, `'stim_meta'` (read fields by key:
+  `batch['stims']`, `batch['responses']`, …). `ccmax` and `ttrc` are no
+  longer dataloader-side pre-computed tensors — they are computed on demand
+  by `normalized_corrcoef` from raw `responses`.
 - The legacy code calls `prediction.squeeze(-2)` to drop the R-axis
   before metrics. The new metrics expect `(B, N, 1, T)` per the model
   paradigm — no squeeze.
diff --git a/docs/_source/md/metrics_paradigm.md b/docs/_source/md/metrics_paradigm.md
index 6c8bac3..4535929 100644
--- a/docs/_source/md/metrics_paradigm.md
+++ b/docs/_source/md/metrics_paradigm.md
@@ -493,7 +493,8 @@ Mirroring `data_paradigm.md` §6, the canonical loop now looks like:
 ```python
 from deepSTRF.metrics import mse_loss, corrcoef, normalized_corrcoef
 
-for stims, responses, valid_mask, stim_metas in loader:
+for batch in loader:                                             # batch is a dict
+    stims, responses = batch['stims'], batch['responses']
     pred = model(stims)                                          # (B, N, 1, T)
 
     loss = mse_loss(pred, responses)                             # auto-PSTH inside
diff --git a/examples/alice_eeg_tutorial.ipynb b/examples/alice_eeg_tutorial.ipynb
index ef95d5a..9f479dc 100644
--- a/examples/alice_eeg_tutorial.ipynb
+++ b/examples/alice_eeg_tutorial.ipynb
@@ -195,8 +195,8 @@
     "    for ep in range(max_epochs):\n",
     "        model.train()\n",
     "        ep_losses = []\n",
-    "        for stims, responses, _, _ in train_loader:\n",
-    "            stims, responses = stims.to(device), responses.to(device)\n",
+    "        for batch in train_loader:\n",
+    "            stims, responses = batch['stims'].to(device), batch['responses'].to(device)\n",
     "            pred = model(stims)\n",
     "            loss = mse_loss(pred, responses)\n",
     "            opt.zero_grad(); loss.backward(); opt.step()\n",
@@ -204,8 +204,8 @@
     "        model.eval()\n",
     "        with torch.no_grad():\n",
     "            v = []\n",
-    "            for stims, responses, _, _ in val_loader:\n",
-    "                stims, responses = stims.to(device), responses.to(device)\n",
+    "            for batch in val_loader:\n",
+    "                stims, responses = batch['stims'].to(device), batch['responses'].to(device)\n",
     "                pred = model(stims); gt = responses.nanmean(dim=2, keepdim=True)\n",
     "                v.append(corrcoef(pred, gt, reduction='none').cpu())\n",
     "            val_cc = torch.stack(v).nanmean().item()\n",
@@ -220,8 +220,8 @@
     "    model.eval()\n",
     "    with torch.no_grad():\n",
     "        cc, ve = [], []\n",
-    "        for stims, responses, _, _ in test_loader:\n",
-    "            stims, responses = stims.to(device), responses.to(device)\n",
+    "        for batch in test_loader:\n",
+    "            stims, responses = batch['stims'].to(device), batch['responses'].to(device)\n",
     "            pred = model(stims); gt = responses.nanmean(dim=2, keepdim=True)\n",
     "            cc.append(corrcoef(pred, gt, reduction='none').cpu())\n",
     "            ve.append(fve(pred, gt, reduction='none').cpu())\n",
diff --git a/examples/crcns_aa_tutorial.ipynb b/examples/crcns_aa_tutorial.ipynb
index 0f7ffdc..53259d4 100644
--- a/examples/crcns_aa_tutorial.ipynb
+++ b/examples/crcns_aa_tutorial.ipynb
@@ -487,7 +487,7 @@
    ],
    "source": [
     "loader = DataLoader(aa1, batch_size=4, shuffle=False, collate_fn=neural_collate)\n",
-    "stims, responses, valid_mask, stim_metas = next(iter(loader))\n",
+    "batch = next(iter(loader))\nstims, responses, valid_mask, stim_metas = batch['stims'], batch['responses'], batch['valid_mask'], batch['stim_meta']\n",
     "\n",
     "print(f\"stims shape       : {tuple(stims.shape)}       (B, 1, F, T_max)\")\n",
     "print(f\"responses shape   : {tuple(responses.shape)}   (B, N_selected, R_max, T_max)\")\n",
@@ -524,7 +524,7 @@
     }
    },
    "outputs": [],
-   "source": "aa2 = CRCNSAA2Dataset(\n    download=True,\n    areas=(\"Field_L\", \"mld\", \"OV\", \"CM\"),\n    stimuli=(\"conspecific\", \"songrip\", \"flatrip\"),\n    dt_ms=DT_MS,\n    smooth=False,\n)\nprint(aa2)\nprint(f\"coverage: {int(aa2.nrn_masks.sum())} / {aa2.nrn_masks.numel()} \"\n      f\"({100 * aa2.nrn_masks.float().mean().item():.1f}%)\")\nprint(f\"one nrn_meta entry: {aa2.nrn_meta[0]}\")\n\n# same API, same 4-tuple from the collate:\naa2.select_population(list(range(min(aa2.N_neurons, 32))))   # take 32 neurons to keep the batch small\nloader = DataLoader(aa2, batch_size=4, shuffle=False, collate_fn=neural_collate)\nstims, responses, valid_mask, stim_metas = next(iter(loader))\nprint(f\"AA2 batch stims shape      : {tuple(stims.shape)}\")\nprint(f\"AA2 batch responses shape  : {tuple(responses.shape)}\")\nprint(f\"AA2 batch valid_mask shape : {tuple(valid_mask.shape)}\")\n"
+   "source": "aa2 = CRCNSAA2Dataset(\n    download=True,\n    areas=(\"Field_L\", \"mld\", \"OV\", \"CM\"),\n    stimuli=(\"conspecific\", \"songrip\", \"flatrip\"),\n    dt_ms=DT_MS,\n    smooth=False,\n)\nprint(aa2)\nprint(f\"coverage: {int(aa2.nrn_masks.sum())} / {aa2.nrn_masks.numel()} \"\n      f\"({100 * aa2.nrn_masks.float().mean().item():.1f}%)\")\nprint(f\"one nrn_meta entry: {aa2.nrn_meta[0]}\")\n\n# same API, same dict from the collate:\naa2.select_population(list(range(min(aa2.N_neurons, 32))))   # take 32 neurons to keep the batch small\nloader = DataLoader(aa2, batch_size=4, shuffle=False, collate_fn=neural_collate)\nbatch = next(iter(loader))\nstims, responses, valid_mask, stim_metas = batch['stims'], batch['responses'], batch['valid_mask'], batch['stim_meta']\nprint(f\"AA2 batch stims shape      : {tuple(stims.shape)}\")\nprint(f\"AA2 batch responses shape  : {tuple(responses.shape)}\")\nprint(f\"AA2 batch valid_mask shape : {tuple(valid_mask.shape)}\")\n"
   },
   {
    "cell_type": "markdown",
diff --git a/examples/dataset_concatenation.ipynb b/examples/dataset_concatenation.ipynb
index 548a70a..6189b89 100644
--- a/examples/dataset_concatenation.ipynb
+++ b/examples/dataset_concatenation.ipynb
@@ -471,7 +471,7 @@
     "combined.select_population(list(range(min(combined.N_neurons, 32))))  # take 32 neurons for a small batch demo\n",
     "loader = DataLoader(combined, batch_size=4, shuffle=False, collate_fn=neural_collate)\n",
     "\n",
-    "stims, responses, valid_mask, stim_metas = next(iter(loader))\n",
+    "batch = next(iter(loader))\nstims, responses, valid_mask, stim_metas = batch['stims'], batch['responses'], batch['valid_mask'], batch['stim_meta']\n",
     "print(f\"stims shape       : {tuple(stims.shape)}\")\n",
     "print(f\"responses shape   : {tuple(responses.shape)}\")\n",
     "print(f\"valid_mask shape  : {tuple(valid_mask.shape)}    (B, N, R, T)\")\n",
diff --git a/examples/explore_nat4.ipynb b/examples/explore_nat4.ipynb
index a9cd2a8..c614d16 100644
--- a/examples/explore_nat4.ipynb
+++ b/examples/explore_nat4.ipynb
@@ -570,7 +570,7 @@
     "print(f\"  visible cells:  {len(ds._selected())}\")\n",
     "\n",
     "# What __getitem__ returns now: only val-data cells, only val stims.\n",
-    "stim, responses_per_neuron, _, meta = ds[0]\n",
+    "item = ds[0]\nstim, responses_per_neuron, meta = item['stims'], item['responses'], item['stim_meta']\n",
     "print(f\"\\nfirst returned item:\")\n",
     "print(f\"  stim shape:  {tuple(stim.shape)}\")\n",
     "print(f\"  N returned:  {len(responses_per_neuron)}\")\n",
@@ -669,7 +669,7 @@
     "ds.select_stims_by_attr('subset', 'val')\n",
     "ds.select_pop_by_nrn_attr('auditory', True)\n",
     "loader = DataLoader(ds, batch_size=4, shuffle=False, collate_fn=neural_collate)\n",
-    "stims, responses, valid_mask, stim_metas = next(iter(loader))\n",
+    "batch = next(iter(loader))\nstims, responses, valid_mask, stim_metas = batch['stims'], batch['responses'], batch['valid_mask'], batch['stim_meta']\n",
     "print(f\"stims:      {tuple(stims.shape)}  (B, 1, F, T)\")\n",
     "print(f\"responses:  {tuple(responses.shape)}  (B, N, R_max, T_max)\")\n",
     "print(f\"valid_mask: {tuple(valid_mask.shape)}  bool\")\n",
diff --git a/tests/test_alice_eeg.py b/tests/test_alice_eeg.py
index 73473ac..bb35e5d 100644
--- a/tests/test_alice_eeg.py
+++ b/tests/test_alice_eeg.py
@@ -134,7 +134,9 @@ def test_alice_dataloader_integration(alice_s01):
     from deepSTRF.utils.data import neural_collate
 
     loader = DataLoader(alice_s01, batch_size=4, collate_fn=neural_collate)
-    stims, responses, valid_mask, stim_metas = next(iter(loader))
+    batch = next(iter(loader))
+    stims, responses, valid_mask, stim_metas = (
+        batch['stims'], batch['responses'], batch['valid_mask'], batch['stim_meta'])
 
     assert stims.dim() == 4 and stims.shape[1:3] == (1, 8)        # (B, 1, F, T)
     assert responses.dim() == 4 and responses.shape[1] == 61      # (B, N, R, T)
diff --git a/tests/test_concat.py b/tests/test_concat.py
index 753116e..ab5a19d 100644
--- a/tests/test_concat.py
+++ b/tests/test_concat.py
@@ -131,12 +131,12 @@ def test_concat_getitem_only_yields_iterable_stims_under_selection():
     # select only A's neurons -> only A's stims iterable
     c.select_population([0, 1])
     assert len(c) == 3
-    seen_metas = [c[i][3] for i in range(len(c))]
+    seen_metas = [c[i]['stim_meta'] for i in range(len(c))]
     assert seen_metas == c.stim_meta[:3]
 
     # every yielded item must have at least one non-NaN response
     for i in range(len(c)):
-        _, resps, _, _ = c[i]
+        resps = c[i]['responses']
         assert any(not r.isnan().any() for r in resps), \
             f"item {i} has fully-NaN responses under selection {c.I}"
 
@@ -147,10 +147,10 @@ def test_concat_getitem_only_yields_iterable_stims_under_selection():
     # select only B's neurons -> only B's stims iterable
     c.select_population([2, 3, 4])
     assert len(c) == 2
-    seen_metas = [c[i][3] for i in range(len(c))]
+    seen_metas = [c[i]['stim_meta'] for i in range(len(c))]
     assert seen_metas == c.stim_meta[3:]
     for i in range(len(c)):
-        _, resps, _, _ = c[i]
+        resps = c[i]['responses']
         assert any(not r.isnan().any() for r in resps)
 
 
@@ -169,11 +169,12 @@ def test_dataloader_over_concat_skips_cross_block_stims():
 
     # 3 iterable stims, batch_size 2 -> 2 batches (sizes 2 and 1)
     assert len(batches) == 2
-    total_items = sum(b_[0].shape[0] for b_ in batches)
+    total_items = sum(b_['stims'].shape[0] for b_ in batches)
     assert total_items == 3
 
     # every batch must contain only A's stims (i.e. valid_mask has at least one True per item)
-    for stims, responses, valid_mask, metas in batches:
+    for b_ in batches:
+        valid_mask = b_['valid_mask']
         per_item_has_data = valid_mask.any(dim=(1, 2, 3))  # (B,) bool
         assert per_item_has_data.all(), \
             "DataLoader yielded a batch item with no valid (s,n,r,t) data anywhere"
diff --git a/tests/test_downer2025_dataset.py b/tests/test_downer2025_dataset.py
index 5fef93d..5fbb56a 100644
--- a/tests/test_downer2025_dataset.py
+++ b/tests/test_downer2025_dataset.py
@@ -283,7 +283,9 @@ def test_timit_collate_produces_correct_shapes(ds_timit_one_session):
     from deepSTRF.utils.data import neural_collate
     loader = DataLoader(ds_timit_one_session, batch_size=2, shuffle=False,
                         collate_fn=neural_collate)
-    stims, resps, mask, metas = next(iter(loader))
+    batch = next(iter(loader))
+    stims, resps, mask, metas = (batch['stims'], batch['responses'],
+                                 batch['valid_mask'], batch['stim_meta'])
     assert stims.shape[:3] == (2, 1, ds_timit_one_session.F)
     assert resps.shape[:2] == (2, ds_timit_one_session.N_neurons)
     assert mask.shape == resps.shape
diff --git a/tests/test_espejo_dataset.py b/tests/test_espejo_dataset.py
index c8d34a9..49cf4c8 100644
--- a/tests/test_espejo_dataset.py
+++ b/tests/test_espejo_dataset.py
@@ -128,7 +128,9 @@ def test_vmn_collate_produces_correct_shapes(vmn_dataset):
 
     loader = DataLoader(vmn_dataset, batch_size=2, shuffle=False,
                         collate_fn=neural_collate)
-    stims, resps, mask, metas = next(iter(loader))
+    batch = next(iter(loader))
+    stims, resps, mask, metas = (batch['stims'], batch['responses'],
+                                 batch['valid_mask'], batch['stim_meta'])
     assert stims.shape[:3] == (2, 1, vmn_dataset.F), stims.shape
     assert resps.shape[:2] == (2, vmn_dataset.N_neurons), resps.shape
     assert mask.shape == resps.shape
diff --git a/tests/test_fitter.py b/tests/test_fitter.py
index 932922a..b7552f9 100644
--- a/tests/test_fitter.py
+++ b/tests/test_fitter.py
@@ -64,7 +64,8 @@ def __getitem__(self, i):
         stim = self.stims[i]
         responses = self.responses_per_stim[i]
         nrn_mask = torch.ones(self.N, dtype=torch.bool)
-        return stim, responses, nrn_mask, {"idx": i}
+        return {'stims': stim, 'responses': responses,
+                'valid_mask': nrn_mask, 'stim_meta': {"idx": i}}
 
 
 class _LinearReadout(torch.nn.Module):
diff --git a/tests/test_ns1_waveform.py b/tests/test_ns1_waveform.py
index 9b1b251..78358a4 100644
--- a/tests/test_ns1_waveform.py
+++ b/tests/test_ns1_waveform.py
@@ -138,7 +138,9 @@ def test_ns1_waveform_collate():
     ds = NS1Dataset(return_waveform=True)
     ds.select_population(list(range(ds.get_N())))
     loader = DataLoader(ds, batch_size=4, collate_fn=neural_collate)
-    stims, responses, valid_mask, metas = next(iter(loader))
+    batch = next(iter(loader))
+    stims, responses, valid_mask, metas = (batch['stims'], batch['responses'],
+                                           batch['valid_mask'], batch['stim_meta'])
     expected_T = 999 * (ds.audio_fs // 200)
     assert stims.shape == (4, 1, expected_T)
     assert responses.shape == (4, ds.get_N(), 20, 999)
diff --git a/tests/test_stim_selection.py b/tests/test_stim_selection.py
index bf4034f..16fe4b8 100644
--- a/tests/test_stim_selection.py
+++ b/tests/test_stim_selection.py
@@ -85,7 +85,7 @@ def test_select_stim_single():
     ds.select_stim(2)
     assert ds.S_sel == [2]
     assert len(ds) == 1
-    _, _, _, meta = ds[0]
+    meta = ds[0]['stim_meta']
     assert meta["name"] == "est2"
 
 
@@ -102,7 +102,7 @@ def test_select_stims_list():
     ds.select_stims([1, 3, 5])
     assert ds.S_sel == [1, 3, 5]
     assert len(ds) == 3
-    metas = [ds[i][3]["name"] for i in range(len(ds))]
+    metas = [ds[i]['stim_meta']["name"] for i in range(len(ds))]
     assert metas == ["est1", "est3", "val1"]
 
 
@@ -113,7 +113,7 @@ def test_select_stims_by_attr():
     assert ds.S_sel == [4, 5]
     assert len(ds) == 2
     for i in range(len(ds)):
-        _, _, _, meta = ds[i]
+        meta = ds[i]['stim_meta']
         assert meta["subset"] == "val"
 
 
@@ -145,7 +145,7 @@ def test_bidirectional_hides_neurons_outside_stim_subset():
     # _selected returns only neurons with at least one valid val response
     assert ds._selected() == [0, 1, 2]
     # __getitem__ batches contain only those neurons
-    _, resps, _, _ = ds[0]
+    resps = ds[0]['responses']
     assert len(resps) == 3
     # none of those responses are NaN sentinels
     for r in resps:
diff --git a/tests/test_wingert2026.py b/tests/test_wingert2026.py
index 703a922..9756dad 100644
--- a/tests/test_wingert2026.py
+++ b/tests/test_wingert2026.py
@@ -417,7 +417,9 @@ def test_collate_roundtrip(clt027c_dataset):
     from torch.utils.data import DataLoader
     ds = clt027c_dataset
     ld = DataLoader(ds, batch_size=4, collate_fn=neural_collate)
-    stims, responses, valid_mask, metas = next(iter(ld))
+    batch = next(iter(ld))
+    stims, responses, valid_mask, metas = (batch['stims'], batch['responses'],
+                                           batch['valid_mask'], batch['stim_meta'])
     B = 4
     assert stims.shape == (B, 1, 32, 2200)
     assert responses.shape[0] == B