fix: pad non-square batches to backbone block size in collator, not in transforms - #992

What does this PR do?

Follow-up to #991. That PR fixed the immediate training crash with square_resize_div_64=False by appending PadIfNeeded(pad_*_divisor=patch_size*num_windows) to the non-square resize pipeline, so the windowed-attention backbone's divisibility assertion passes. That works, but it leaves a subtle residual issue with the NestedTensor mask.

This PR moves the divisibility padding from the transform pipeline into the batch collator. Every zero-pad pixel (batch-level AND divisor round-up) is now correctly marked as padding in the mask, so downstream attention ignores it instead of treating it as real content.

Why the previous approach under-counts pad in the mask

PadIfNeeded inside the transform pipeline pads each image before it reaches the collator. The collator then builds each image's mask from the per-image shape after transforms, so the padded region is reported as "content" even though it's actually zeros.

Concrete trace: mixed landscape + portrait batch

Config: resolution=896, max_size=1333, patch_size=16, num_windows=2, so block_size=32.

Input batch (H × W):

• Image A (landscape): 800 × 1200
• Image B (portrait): 1200 × 800

Each image through the current pipeline:

Now the batch collator:

• batch_max_H = max(896, 1344) = 1344
• batch_max_W = max(1344, 896) = 1344
• Batch tensor shape: (2, 3, 1344, 1344)

Per-image pad breakdown inside the 1344 × 1344 slot:

Image A (landscape, effective shape 896 × 1344):

• Bottom 448 rows: batch-level pad, mask marks pad (correct).
• 7 extra rows (889 to 896) and 11 cols (1333 to 1344) added by PadIfNeeded: inside the mask's "content" region, so mask marks them as content (wrong, they are actually zeros).

Image B (portrait, effective shape 1344 × 896):

• Right 448 cols: batch-level pad, mask marks pad (correct).
• 11 extra rows (1333 to 1344) and 7 cols (889 to 896) added by PadIfNeeded: inside the mask's "content" region, so mask marks them as content (wrong).

After this PR, the transform-level PadIfNeeded is gone. Per-image outputs are 889 × 1333 and 1333 × 889 (natural post-resize shapes). The collator computes batch_max = (1333, 1333) and rounds up to (1344, 1344). Every zero-pad cell (batch-level AND the divisor round-up) is marked pad in the mask. Downstream attention correctly ignores all of it.

Fix

1. Add an optional block_size parameter to nested_tensor_from_tensor_list (and its ONNX variant). When set, round max_size[1]/max_size[2] up to the next multiple of block_size before allocating the batch tensor. The rounded-up strip is explicitly marked True in the mask.
2. Add a picklable make_collate_fn(block_size) factory (uses functools.partial over a module-level helper so it survives multi-process DataLoaders and DDP spawn).
3. In RFDETRDataModule.__init__, compute block_size = model_config.patch_size * model_config.num_windows once and use the factory at every DataLoader site.
4. Drop the per-image PadIfNeeded, the divisor kwarg on _build_train_resize_config, and the _pad_to_divisor_config helper. Backbone still receives divisibility-compliant input, and the mask is now correct.

Related Issue(s): Follow-up to #991. The original training crash it fixed was issue #983.

Type of Change

• Bug fix (non-breaking change that fixes an issue)

Testing

• I have tested this change locally
• I have added/updated tests for this change

Test details:

• Added TestNestedTensorBlockSize (5 tests): verify block_size=None preserves old behavior, block_size=32 rounds batch-max up, parametrized for several block sizes (32, 56, 64), and that every pad cell (content vs divisor strip) is correctly marked in the mask.
• Added TestMakeCollateFn (4 tests): verify the factory's default is backward-compatible, rounds batch-max correctly with block_size=32, passes targets through unchanged, and that mixed landscape+portrait batches have every pad cell correctly masked.
• Removed TestBuildTrainResizeConfigDivisor and TestNonSquareResizeDivisibilityRegression from tests/datasets/test_coco_resize_config.py (those pinned the removed divisor kwarg / PadIfNeeded step).
• Reverted two count assertions in tests/datasets/test_augmentations.py: val/test pipeline now has 2 resize wrappers (SmallestMaxSize + LongestMaxSize) instead of 3.
• 951 tests in (tests/datasets/, tests/utilities/, tests/training/) all pass locally.
• pre-commit run clean on every touched file.

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code where necessary, particularly in hard-to-understand areas
• My changes generate no new warnings or errors
• I have updated the documentation accordingly (if applicable)