fix(nemotron_omni): support mixed-resolution image batches in tile path (#3559)
`nemotron_omni_collate_fn` cannot stack images of differing resolutions, so it returns `pixel_values` as a list of per-image tensors. The tile-based branch of `forward()` did `vit_batch_size = pixel_values.shape[0]` unconditionally, raising `AttributeError: 'list' object has no attribute 'shape'`. This made `local_batch_size > 1` unusable for image SFT on any variable-resolution dataset (e.g. CORD-V2). The collate comment claimed `extract_feature` already handled list input, but it did not: it called the vision tower and read `.shape` directly, and HF RADIO only accepts a dense (B, C, H, W) tensor. - `extract_feature` now dispatches on input type; the dense path is factored into `_extract_feature_dense` and run once per image for list input. Eval-mode toggling moves into try/finally so it survives the loop. - `forward()` derives `vit_batch_size` from `len()` for list input. - The `image_flags` filter indexed a stacked [num_images, num_tokens, C] tensor, which cannot exist when token counts vary by resolution. The list path now filters per image and concatenates along the token dim. Adds CPU-only unit tests covering a mixed-resolution batch, per-image `image_flags` filtering, train-mode restoration, and the unchanged dense path. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
H
Huiying committed
8b4db0075182b2f9d82e9c8f9342d911aea9ed6c
Parent: 79fd3a4
Committed by GitHub <noreply@github.com>
on 8/15/2026, 12:46:46 AM