diff --git a/giant/cli.py b/giant/cli.py
index ea254bb..6b6bab2 100644
--- a/giant/cli.py
+++ b/giant/cli.py
@@ -977,7 +977,7 @@ def predict(
         nonlocal writer, total
 
         if coord == Coord.local:
-            cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(
+            feats = build_features(
                 piece,
                 pdg_map,
                 mat_map,
@@ -987,7 +987,9 @@ def predict(
                 mat_topn_map=cond_mat_topn,
                 k_max=stage2_k_max,
             )
-            cond_cont = cond_norm.transform(cond_cont)
+            cond_cat = feats.cond_cat
+            target_raw = feats.target_s1
+            cond_cont = cond_norm.transform(feats.cond_cont)
         else:
             cond_cont, cond_cat = build_cond_features(
                 piece,
diff --git a/giant/data/dataset.py b/giant/data/dataset.py
index 6e2e282..a82ad4b 100644
--- a/giant/data/dataset.py
+++ b/giant/data/dataset.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import NamedTuple
 
 import numpy as np
 import torch
@@ -11,6 +12,37 @@ from giant.data.loader import event_id_offset, iter_file_chunks
 from giant.data.transforms import Normalizer, build_features, sorted_membership
 
 
+class StepBatch(NamedTuple):
+    """One training batch, as yielded by `StreamingStepsDataset`. Field order
+    is load-bearing for existing positional unpacking elsewhere (`trainers.py`,
+    `validate.py`, test fixtures) — append only, never insert or reorder.
+
+    cond_cont:  (B, COND_DIM) float32
+    cond_cat:   (B, 2/3/4) int64 — width 2 unless conditioning="onehot"
+    target_s1:  (B, 9) float32 — normalised Stage-1 primary target
+    n_sec:      (B,) int64 — true secondary count per step
+    sec_cont:   (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
+                local_dir, log_mass, charge] per slot (mass/charge
+                normalised iff `sec_phys_normalizer` was given); always
+                computed the same way regardless of
+                stage2_model.particle_type.target, only actually used
+                downstream under target="physical"
+    proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
+                only; zeros when `proc_map` is None)
+    sec_type_idx: (B, k_max) int64 — per-slot class index into
+                `sec_type_class_map`, for particle_type.target in
+                ("onehot", "embedding"); zeros (unused) otherwise
+    """
+
+    cond_cont: torch.Tensor
+    cond_cat: torch.Tensor
+    target_s1: torch.Tensor
+    n_sec: torch.Tensor
+    sec_cont: torch.Tensor
+    proc_idx: torch.Tensor
+    sec_type_idx: torch.Tensor
+
+
 def make_event_split(
     all_event_ids: np.ndarray,
     val_fraction: float = 0.1,
@@ -39,24 +71,7 @@ class StreamingStepsDataset(IterableDataset):
     rather than single rows, so the batch is assembled with vectorized
     numpy slicing instead of a per-row Python loop in the default collate.
 
-    Each batch is a tuple:
-        (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)
-    where:
-        cond_cont:  (B, COND_DIM) float32
-        cond_cat:   (B, 2/3/4) int64 — width 2 unless conditioning="onehot"
-        target_s1:  (B, 9) float32 — normalised Stage-1 primary target
-        n_sec:      (B,) int64 — true secondary count per step
-        sec_cont:   (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
-                    local_dir, log_mass, charge] per slot (mass/charge
-                    normalised iff `sec_phys_normalizer` was given); always
-                    computed the same way regardless of
-                    stage2_model.particle_type.target, only actually used
-                    downstream under target="physical"
-        proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
-                    only; zeros when `proc_map` is None)
-        sec_type_idx: (B, k_max) int64 — per-slot class index into
-                    `sec_type_class_map`, for particle_type.target in
-                    ("onehot", "embedding"); zeros (unused) otherwise
+    Each batch is a `StepBatch` — see its docstring for field meanings.
 
     `k_max` (constructor arg, default the module constant) should match
     `stage2_model.k_max` — it sets the padded
@@ -129,17 +144,7 @@ class StreamingStepsDataset(IterableDataset):
                     continue
                 chunk = {k: v[mask] for k, v in chunk.items()}
 
-                (
-                    cond_cont,
-                    cond_cat,
-                    target_s1,
-                    n_sec,
-                    sec_cont,
-                    proc_idx,
-                    sec_type_idx,
-                    _,
-                    _,
-                ) = build_features(
+                feats = build_features(
                     chunk,
                     self.pdg_map,
                     self.mat_map,
@@ -155,14 +160,14 @@ class StreamingStepsDataset(IterableDataset):
                     sec_type_class_map=self.sec_type_class_map,
                     k_max=self.k_max,
                 )
-                buf_cont.append(cond_cont)
-                buf_cat.append(cond_cat)
-                buf_tgt.append(target_s1)
-                buf_nsec.append(n_sec)
-                buf_sec.append(sec_cont)
-                buf_proc.append(proc_idx)
-                buf_type.append(sec_type_idx)
-                buf_n += len(cond_cont)
+                buf_cont.append(feats.cond_cont)
+                buf_cat.append(feats.cond_cat)
+                buf_tgt.append(feats.target_s1)
+                buf_nsec.append(feats.n_sec)
+                buf_sec.append(feats.sec_cont)
+                buf_proc.append(feats.proc_idx)
+                buf_type.append(feats.sec_type_idx)
+                buf_n += len(feats.cond_cont)
 
                 if buf_n >= self.shuffle_buffer:
                     (
@@ -226,14 +231,14 @@ class StreamingStepsDataset(IterableDataset):
         n_full = n // bs if not final else (n + bs - 1) // bs
         for start in range(0, n_full * bs, bs):
             end = min(start + bs, n)
-            yield (
-                torch.from_numpy(cont[start:end]).float(),
-                torch.from_numpy(cat[start:end]).long(),
-                torch.from_numpy(tgt[start:end]).float(),
-                torch.from_numpy(nsec[start:end]).long(),
-                torch.from_numpy(sec[start:end]).float(),
-                torch.from_numpy(proc[start:end]).long(),
-                torch.from_numpy(styp[start:end]).long(),
+            yield StepBatch(
+                cond_cont=torch.from_numpy(cont[start:end]).float(),
+                cond_cat=torch.from_numpy(cat[start:end]).long(),
+                target_s1=torch.from_numpy(tgt[start:end]).float(),
+                n_sec=torch.from_numpy(nsec[start:end]).long(),
+                sec_cont=torch.from_numpy(sec[start:end]).float(),
+                proc_idx=torch.from_numpy(proc[start:end]).long(),
+                sec_type_idx=torch.from_numpy(styp[start:end]).long(),
             )
 
         if final:
diff --git a/giant/data/transforms.py b/giant/data/transforms.py
index 4dfc672..5f36fcf 100644
--- a/giant/data/transforms.py
+++ b/giant/data/transforms.py
@@ -1,4 +1,5 @@
 import warnings
+from typing import NamedTuple
 
 import numpy as np
 
@@ -860,36 +861,10 @@ def _cond_normalizer_transform(
     return ((cond_cont - mean) / std).astype(np.float32)
 
 
-def build_features(
-    data: dict[str, np.ndarray],
-    pdg_map: dict[int, int],
-    mat_map: dict[str, int],
-    cond_normalizer: Normalizer | None = None,
-    target_normalizer: Normalizer | None = None,
-    sec_phys_normalizer: Normalizer | None = None,
-    fit: bool = False,
-    proc_map: dict[str, int] | None = None,
-    require_secondaries: bool = False,
-    particle_conditioning: str = "embedding",
-    material_conditioning: str = "embedding",
-    sec_phys_only: bool = False,
-    pdg_topn_map: dict[int, int] | None = None,
-    mat_topn_map: dict[str, int] | None = None,
-    sec_type_class_map: dict | None = None,
-    k_max: int = K_MAX,
-) -> tuple[
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    Normalizer | None,
-    Normalizer | None,
-]:
-    """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
-    sec_type_idx) arrays.
+class StepFeatures(NamedTuple):
+    """Output of `build_features`. Field order is load-bearing for existing
+    positional unpacking (tests, `StreamingStepsDataset`) — append only,
+    never insert or reorder.
 
     target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
     n_sec: (N,) integer secondary counts (target for n_sec head)
@@ -908,6 +883,40 @@ def build_features(
               in `("onehot", "embedding")` — see `encode_secondary_type_idx`.
               Zero-filled (and unused) when `sec_type_class_map` is None
               (i.e. `target = "physical"`).
+    """
+
+    cond_cont: np.ndarray
+    cond_cat: np.ndarray
+    target_s1: np.ndarray
+    n_sec: np.ndarray
+    sec_cont: np.ndarray
+    proc_idx: np.ndarray
+    sec_type_idx: np.ndarray
+    cond_normalizer: Normalizer | None
+    target_normalizer: Normalizer | None
+
+
+def build_features(
+    data: dict[str, np.ndarray],
+    pdg_map: dict[int, int],
+    mat_map: dict[str, int],
+    cond_normalizer: Normalizer | None = None,
+    target_normalizer: Normalizer | None = None,
+    sec_phys_normalizer: Normalizer | None = None,
+    fit: bool = False,
+    proc_map: dict[str, int] | None = None,
+    require_secondaries: bool = False,
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
+    sec_phys_only: bool = False,
+    pdg_topn_map: dict[int, int] | None = None,
+    mat_topn_map: dict[str, int] | None = None,
+    sec_type_class_map: dict | None = None,
+    k_max: int = K_MAX,
+) -> StepFeatures:
+    """Assemble a `StepFeatures` of (cond_cont, cond_cat, target_s1, n_sec,
+    sec_cont, proc_idx, sec_type_idx, cond_normalizer, target_normalizer) —
+    see `StepFeatures` for field meanings.
 
     require_secondaries: when True, raise if any step has n_sec > 0 but the
         per-secondary list columns are absent (a mis-converted file that would
@@ -1054,14 +1063,14 @@ def build_features(
     else:
         proc_idx = np.zeros(len(cond_cat), dtype=np.int64)
 
-    return (
-        cond_cont,
-        cond_cat,
-        target_s1,
-        n_sec,
-        sec_cont,
-        proc_idx,
-        sec_type_idx,
-        cond_normalizer,
-        target_normalizer,
+    return StepFeatures(
+        cond_cont=cond_cont,
+        cond_cat=cond_cat,
+        target_s1=target_s1,
+        n_sec=n_sec,
+        sec_cont=sec_cont,
+        proc_idx=proc_idx,
+        sec_type_idx=sec_type_idx,
+        cond_normalizer=cond_normalizer,
+        target_normalizer=target_normalizer,
     )
diff --git a/giant/pipeline.py b/giant/pipeline.py
index 683053a..e68a9de 100644
--- a/giant/pipeline.py
+++ b/giant/pipeline.py
@@ -248,7 +248,7 @@ def run_setup_stage(
                 if not mask.any():
                     continue
                 chunk_tr = {k: v[mask] for k, v in chunk.items()}
-                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _, _ = build_features(
+                feats = build_features(
                     chunk_tr,
                     pdg_map,
                     mat_map,
@@ -259,6 +259,10 @@ def run_setup_stage(
                     sec_phys_only=True,
                     k_max=k_max,
                 )
+                cond_cont = feats.cond_cont
+                target_s1 = feats.target_s1
+                n_sec = feats.n_sec
+                sec_cont = feats.sec_cont
                 cond_acc.update(cond_cont)
                 tgt_acc.update(target_s1)
                 if energy_sampler is not None:
diff --git a/giant/training/trainers.py b/giant/training/trainers.py
index d6ffd6d..f3ba4a4 100644
--- a/giant/training/trainers.py
+++ b/giant/training/trainers.py
@@ -16,6 +16,7 @@ adversarial and non-adversarial stages identically.
 import copy
 import math
 from dataclasses import dataclass, field
+from typing import NamedTuple
 
 import torch
 import torch.nn.functional as F
@@ -23,6 +24,7 @@ import torch.optim as optim
 
 from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
 from giant.constants import CONT_SLOT_DIM
+from giant.data.dataset import StepBatch
 from giant.model.network import Router, stage2_type_dim
 from giant.model.schedule import (
     CosineSchedule,
@@ -72,8 +74,8 @@ def _cosine_warmup_lambda(warmup_steps: int, total_steps: int):
     return _lr_lambda
 
 
-def _batch_to_device(batch: tuple, device: torch.device) -> tuple:
-    return tuple(t.to(device) for t in batch)
+def _batch_to_device(batch: StepBatch, device: torch.device) -> StepBatch:
+    return type(batch)(*(t.to(device) for t in batch))
 
 
 @dataclass(frozen=True)
@@ -185,11 +187,10 @@ class StageSpec:
 class StageTrainer:
     """One active stage's optimizer(s), EMA, and per-batch step.
 
-    Reads only the shared batch tuple `(cond_cont, cond_cat, x1_s1, n_sec,
-    sec_cont, proc_idx, sec_type_idx)` — stage 2 always conditions on the
-    ground-truth `x1_s1` (`stage2_model.stage1_context = "truth"`,
-    stage-level teacher forcing; `"sampled"` is not implemented), so stage
-    trainers never need each other's output at train time. This means
+    Reads only the shared `StepBatch` (`giant.data.dataset`) — stage 2 always
+    conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context =
+    "truth"`, stage-level teacher forcing; `"sampled"` is not implemented),
+    so stage trainers never need each other's output at train time. This means
     "stage-2-only training is a cheap ablation, not new plumbing" falls out
     for free: a trainer only exists for active stages, and inactive stages
     are simply never constructed.
@@ -255,10 +256,10 @@ class StageTrainer:
 
     # --- per-batch (subclass responsibility) ----------------------------
 
-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
         raise NotImplementedError
 
-    def val_loss(self, batch: tuple, device: torch.device) -> dict:
+    def val_loss(self, batch: StepBatch, device: torch.device) -> dict:
         raise NotImplementedError
 
     # --- reporting hooks ------------------------------------------------
@@ -575,7 +576,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             l_type = (se * mask).sum() / denom
         return l_type, type_acc
 
-    def _compute(self, batch: tuple, device: torch.device, epoch: int | None = None) -> dict:
+    def _compute(self, batch: StepBatch, device: torch.device, epoch: int | None = None) -> dict:
         """`epoch=None` (the `val_loss` path) always uses full teacher
         forcing (`p_tf=1.0`) regardless of `spec.teacher_forcing` — validation
         should stay a stable, non-stochastic ground-truth comparison; only
@@ -639,7 +640,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             "nsec_acc": nsec_acc,
         }
 
-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
         if self.router is not None:
             self.router.gumbel_tau = _gumbel_tau(
                 global_step,
@@ -659,7 +660,7 @@ class FlowDDPMStageTrainer(StageTrainer):
         return stats
 
     @torch.no_grad()
-    def val_loss(self, batch: tuple, device: torch.device) -> dict:
+    def val_loss(self, batch: StepBatch, device: torch.device) -> dict:
         return {key: value.item() for key, value in self._compute(batch, device).items()}
 
     # --- reporting ------------------------------------------------------
@@ -674,6 +675,16 @@ class FlowDDPMStageTrainer(StageTrainer):
         return val_means.get("loss", 0.0)
 
 
+class _Stage2RealFakeBatch(NamedTuple):
+    """Subset of `StepBatch` that `_stage2_real_and_fake` needs."""
+
+    cond_cont: torch.Tensor
+    cond_cat: torch.Tensor
+    n_sec: torch.Tensor
+    sec_cont: torch.Tensor
+    sec_type_idx: torch.Tensor
+
+
 class WGANStageTrainer(StageTrainer):
     """WGAN-GP generator+critic for a single stage (see giant/model/wgan.py).
 
@@ -737,7 +748,7 @@ class WGANStageTrainer(StageTrainer):
         self.val_metrics = []
         self.stage_metrics = [stage_metric("lr"), stage_metric("critic_lr")]
 
-    def _stage2_real_and_fake(self, batch_tensors, stage1_ctx, global_step, device):
+    def _stage2_real_and_fake(self, batch_tensors: _Stage2RealFakeBatch, stage1_ctx, global_step, device):
         """Build `(real, fake_raw, mask, critic_fn)` for stage 2, covering
         both decoders and all three particle-type targets. `fake_raw` still
         needs the caller's straight-through relaxation under
@@ -777,7 +788,7 @@ class WGANStageTrainer(StageTrainer):
 
         return real, fake_raw, mask, critic_fn
 
-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
         (
             cond_cont,
             cond_cat,
@@ -802,7 +813,7 @@ class WGANStageTrainer(StageTrainer):
             mask = None
         else:
             real, fake_raw, mask, critic_fn = self._stage2_real_and_fake(
-                (cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
+                _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
                 stage1_ctx,
                 global_step,
                 device,
diff --git a/giant/validate.py b/giant/validate.py
index 1d3b5a0..acf3f87 100644
--- a/giant/validate.py
+++ b/giant/validate.py
@@ -115,11 +115,10 @@ def validate_marginals(
     for i, batch in enumerate(val_loader):
         if n_batches is not None and i >= n_batches:
             break
-        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
-        # sec_type_idx).
-        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx, sec_type_idx = batch
-        cond_cont = cond_cont.to(device)
-        cond_cat = cond_cat.to(device)
+        # batch is a StepBatch (giant.data.dataset).
+        x1, n_sec, sec_cont, sec_type_idx = batch.target_s1, batch.n_sec, batch.sec_cont, batch.sec_type_idx
+        cond_cont = batch.cond_cont.to(device)
+        cond_cat = batch.cond_cat.to(device)
 
         gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps)
 
diff --git a/issues.md b/issues.md
index 0c8f62a..c580f98 100644
--- a/issues.md
+++ b/issues.md
@@ -46,7 +46,7 @@ architecture matrix grows, not about rot or breakage.
 | 4 | `cli.py` is at 35.8 % coverage and holds untested override-precedence logic | **High** | Medium | **Fixed** (`2bfb1ab`, partial — see status note) |
 | 5 | Inference bootstrap is duplicated verbatim between `predict` and `rollout` | **High** | Small | **Fixed** |
 | 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | **Fixed** |
-| 7 | Positional tuple contracts (9-tuple, 7-tuple) between data, model and training layers | Medium | Small | Open |
+| 7 | Positional tuple contracts (9-tuple, 7-tuple) between data, model and training layers | Medium | Small | **Fixed** |
 | 8 | `network.py` is 1745 lines holding three distinct modules | Medium | Small | Open |
 | 9 | `scripts` is published as a top-level distribution package | Medium | Small | Open |
 | 10 | `torch.load(weights_only=False)` — checkpoints are arbitrary pickles | Low | Medium | Open |
@@ -891,6 +891,39 @@ reasoning, two hand-maintained copies of a ~10-line error message
 
 ## Issue 7 — Positional tuple contracts between the data, model and training layers
 
+> **Status: Fixed.** `giant/data/transforms.py` now declares `StepFeatures(NamedTuple)`
+> (same 9 fields, same order) and `build_features` returns it instead of a bare tuple;
+> `giant/data/dataset.py` declares `StepBatch(NamedTuple)` (same 7 fields, same order) and
+> `StreamingStepsDataset` yields it. `_batch_to_device` (`trainers.py`) is retyped
+> `(batch: StepBatch, device) -> StepBatch` and reconstructs via `type(batch)(*(...))` so
+> the concrete `NamedTuple` type — not just a plain tuple — survives `.to(device)`,
+> matching what torch's own `pin_memory` already does for `NamedTuple`s (verified against
+> the installed torch 2.3.1: NamedTuple batches survive `DataLoader` worker-process
+> pickling and are reconstructed as the same type, whereas a plain tuple is silently
+> downgraded to a `list` by `pin_memory`). Investigation while fixing this corrected two
+> details in the original issue text: `_batch_to_device` only ever has one real shape (the
+> 7-tuple, used identically at both call sites) — the "5-tuple at a second
+> `_batch_to_device` call site" was a misreading, that 5-element tuple is a **derived**
+> bundle built from 5 of the same 7 already-`_batch_to_device`'d fields
+> (`WGANStageTrainer._stage2_real_and_fake`), now its own `_Stage2RealFakeBatch(NamedTuple)`
+> (trainers-internal); and `giant/validate.py:120` turned out to be a third real consumer
+> of the same contract the issue didn't list, converted to `StepBatch` attribute access.
+> The already-fully-named positional unpacks in `trainers.py`'s `_compute`/`step` were
+> left as positional tuple-unpacks (still valid against a `NamedTuple`, per the issue's own
+> "migration incremental" framing) — only the sites with heavy `_`-throwaways or bare
+> unlabelled unpacks (`dataset.py`'s `build_features` consumption, `cli.py:980`,
+> `pipeline.py:251`, `validate.py:120`) were converted to attribute access. Two test
+> batch-construction helpers (`tests/test_train.py::_fake_batches`,
+> `tests/test_validate.py::_loader`) now build real `StepBatch`s instead of plain tuples,
+> which is what let `_batch_to_device` stay a simple, non-defensive one-liner. Verified live
+> (then reverted) that inserting a field into `StepBatch` now makes `ty check` fail with
+> "Too many values to unpack" at every positional-unpack call site and
+> "No argument provided for required parameter" at every keyword-constructed one — exactly
+> the protection this issue asked for, and precisely the failure mode that was previously
+> invisible. No field was reordered anywhere. `uv run pytest -q` (804 passed), `ruff check`,
+> `ruff format --check`, and `ty check` all clean. Everything below this point describes the
+> pre-fix state and is kept for historical context.
+
 **Severity: Medium. Effort: Small.**
 
 **Location:** `giant/data/transforms.py:863-890` (`build_features`),
diff --git a/tests/test_train.py b/tests/test_train.py
index fc5781d..ef0787f 100644
--- a/tests/test_train.py
+++ b/tests/test_train.py
@@ -17,6 +17,7 @@ from giant.constants import (
     SEC_SLOT_DIM,
     X_DIM,
 )
+from giant.data.dataset import StepBatch
 from giant.model.network import build_critics, build_models
 from giant.training import (
     FlowDDPMStageTrainer,
@@ -316,7 +317,7 @@ def _fake_batches(n_batches, batch_size, seed=0):
         sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
         proc_idx = torch.zeros(batch_size, dtype=torch.long)
         sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
-        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
+        batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches
 
 
diff --git a/tests/test_validate.py b/tests/test_validate.py
index ca220ae..d19a80f 100644
--- a/tests/test_validate.py
+++ b/tests/test_validate.py
@@ -2,6 +2,7 @@ import numpy as np
 import torch
 
 from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
+from giant.data.dataset import StepBatch
 from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim
 from giant.validate import validate_marginals
 
@@ -46,8 +47,7 @@ def _tiny_models(particle_type_cfg: dict | None = None):
 
 
 def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8):
-    """A val_loader matching StreamingStepsDataset's 7-tuple batch shape:
-    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)."""
+    """A val_loader matching StreamingStepsDataset's StepBatch shape."""
     batches = []
     for _ in range(n_batches):
         cond_cont = torch.randn(B, COND_DIM)
@@ -57,7 +57,7 @@ def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int
         sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
         proc_idx = torch.zeros(B, dtype=torch.long)
         sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
-        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
+        batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches