SIGN IN SIGN UP

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