[ChrisW09]

Fixes #577

Problem

The training loader keeps a trailing partial batch. When the number of training rows leaves a remainder of 1 for batch_size, that batch has a single row, and every nn.BatchNorm1d layer fails in training mode with ValueError: Expected more than 1 value per channel when training. With the defaults that is about one dataset size in 128 (e.g. 162 rows → 129 training rows), for models that use torch.nn.BatchNorm1d, e.g. MLPConfig(batch_norm=True) or head_use_batch_norm=True.

Fix

TabularDataModule.train_dataloader() (deeptab/data/datamodule.py) sets drop_last=True for the training loader only when:

• the user hasn't set drop_last in dataloader_kwargs, and
• the training set is larger than one batch, and
• the last batch would contain exactly one row.

Under distributed training the check uses the rows each process receives: Lightning's distributed sampler gives every process ceil(n / world_size) rows, and the global count could otherwise drop a process's only batch (so it trains nothing) or miss a per-process one-row tail.

With shuffling, a different row is left out in each epoch. An explicit dataloader_kwargs={"drop_last": ...} always wins, dataloader_kwargs itself is not modified, and validation, test and prediction loaders are unchanged. Invalid batch sizes (e.g. 0) are not inspected here, so DataLoader still rejects them with its own ValueError (covered by tests/lifecycle/test_continuation.py).

Contrastive pretraining trains on the same loader, so this also removes the training-side one-row batch reported in #529. #530 handles batches that are too small inside the pretraining steps, which also covers validation batches; the two changes are independent.

Tests

• tests/data/test_datamodule.py::test_single_row_trailing_training_batch_is_dropped[...] covers a one-row tail (dropped), a two-row tail and no tail (kept), a single-row dataset (kept), and an explicit drop_last=False (kept).
• tests/data/test_datamodule.py::test_single_row_tail_check_uses_rows_per_process[...] stubs a two-process trainer: 66 training rows (33 per process, one-row tail) are dropped, 33 rows (17 per process) are kept.
• tests/integration/test_partial_training_batches.py (new) checks that MLPRegressor(MLPConfig(batch_norm=True)) and a Mambular head with head_use_batch_norm=True fit and predict with 129 training rows.

The one-row-tail and both estimator cases fail on main and pass with this change. I also ran a real two-process CPU DDP fit: with 162 rows each process trains one 65-row batch, as on main, and with 323 rows and BatchNorm each process trains one 128-row batch, where main crashes.

Docs

• The "Partial batches" section of docs/core_concepts/training_and_evaluation.md describes the automatic drop and how to keep the batch.
• The TabularDataModule and train_dataloader docstrings mention it.

Checks

• just test: 3245 passed, 46 skipped, 28 xfailed; coverage 91.32%
• just docs: build succeeded (warnings treated as errors)
• just check: all hooks passed (ruff, prettier, pyright)