Batch StreamingStepsDataset internally instead of per-row collate

The dataset yielded one row at a time, forcing DataLoader's default
collate to Python-loop over every row to assemble each batch. That
loop scales with batch size and was pinning a CPU core at 100% while
the GPU sat idle. Now the dataset yields whole batches via vectorized
numpy slicing, used with DataLoader(batch_size=None).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

scripts/train.py

@@ -154,6 +154,7 @@ def main() -> None:
         mat_map=mat_map,
         cond_normalizer=cond_norm,
         target_normalizer=tgt_norm,
+        batch_size=t["batch_size"],
         shuffle_buffer=args.shuffle_buffer,
         shuffle=True,
     )
@@ -164,16 +165,19 @@ def main() -> None:
         mat_map=mat_map,
         cond_normalizer=cond_norm,
         target_normalizer=tgt_norm,
+        batch_size=t["batch_size"],
         shuffle=False,
     )
 
+    # Dataset yields whole batches already, so batch_size=None tells DataLoader
+    # to pass them through instead of re-collating row-by-row in Python.
     pin = device.type == "cuda"
     train_loader = DataLoader(
-        train_ds, batch_size=t["batch_size"],
+        train_ds, batch_size=None,
         num_workers=t["num_workers"], pin_memory=pin,
     )
     val_loader = DataLoader(
-        val_ds, batch_size=t["batch_size"],
+        val_ds, batch_size=None,
         num_workers=t["num_workers"], pin_memory=pin,
     )