fix(models): avoid batch-wide float target masks on CUDA

[JESUSROYETH]

CUDA mask-label sampling currently concatenates the matched masks across the batch, then converts the full tensor to float32. In the workload below, SegSmall at batch 20 ran out of memory at that conversion on an L4.

Changes

The fix samples one image at a time with the same native nearest-neighbour sampler, then concatenates the sampled labels. In the loss's no-grad path, each image's full float temporary can be released before sampling the next. CPU and unsupported CUDA inputs keep the existing path; no new option or change to the loss formulas.

Performance

I measured this with RFDETRSegSmall/RFDETRSegMedium, PyTorch 2.9.1+cu129, on a COCO val2017 subset (512 training / 16 validation images), one epoch, seed 51, two workers, default multi-scale and EMA. Times cover the public .train() call end to end, including validation and checkpoints, and repeated values are medians.

Model / batch / augmentation	Runs per arm	Peak allocated GiB, before → after	Train seconds, before → after
Small / 4 / Kornia	3	5.46 → 4.96	53.80 → 53.09
Small / 8 / CPU torchvision	2	9.90 → 8.96	45.49 → 45.10
Small / 16 / Kornia	1	16.54 → 15.30	46.29 → 46.38
Small / 20 / Kornia	2	OOM → 19.63	OOM → 48.78
Medium / 4 / Kornia	1	7.18 → 6.83	69.94 → 68.91

Batch 20 failed twice before the change and completed twice after. The original also failed with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Timing ranges overlap at batch 4 (52.59–53.86 / 52.25–53.14 s) and batch 8 (44.98–46.01 / 44.96–45.23 s), so this is a memory/capacity fix, not a speedup claim .. it doesn't establish the maximum supported batch size either.

Validation

The criterion suite passed 60 tests (44 CPU and 16 CUDA, one XLA skip), plus 3 doctests. Pre-commit is clean, strict mypy included. I checked the sampled labels for byte-identity in 20 real-training sampler calls across Kornia/torchvision, with baseline-self and active-route checks. Full-training baseline repeats already give different final weights on their own, so byte parity here is about the sampled labels, not the final model. DDP and compiled training were not benchmarked.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87%. Comparing base (f0fc53c) to head (6a71fb9).

Additional details and impacted files

@@           Coverage Diff           @@
##           develop   #1437   +/-   ##
=======================================
  Coverage       87%     87%           
=======================================
  Files          114     114           
  Lines        15099   15108    +9     
=======================================
+ Hits         13062   13071    +9     
  Misses        2037    2037           

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.