Repository navigation
[NemotronH-Omni] Fix device mismatch in test tensor creation - #49102
Conversation
`image_grid_hw`, `positions`, and `input_features_mask` were created without `device=torch_device`, causing RuntimeError when tests run on GPU: CPU tensors collided with GPU model buffers (e.g. `summary_idxs` in RadioModel._forward_packed, `lengths` in create_audio_mask). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
|
run-slow: nemotron_h_omni |
Nvidia CIThis comment contains models: ["models/nemotron_h_omni"] |
AMD CIThis comment contains models: ["models/nemotron_h_omni"] |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
…ata-dependent shapes)
|
[For maintainers] Suggested jobs to run (before merge) run-slow: nemotron_h_omni |
CI recapDashboard: View test results in Grafana |
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Yep totally reasonable with moving to device but I swear the run slow also didnt have that many failures on nemotron omni (i.e. mostly export failures)
What does this PR fix?
Fixes 500+ test failures in
NemotronHOmniVision2TextModelTestandNemotronHOmniAudio2TextModelTestintroduced by #46509 (merged 2026-09-24).Root cause
Four tensors in the test file were created without
device=torch_device, causingRuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!on GPU runners:NemotronHOmniVision2TextModelTester._prepare_modality_inputsimage_grid_hwsummary_idxsbuffer inRadioModel._forward_packedNemotronHOmniAudio2TextModelTester.create_audio_maskpositions(arange)lengthsfromids_tensor(on GPU)NemotronHOmniModelTestMixin._image_features_prepare_config_and_inputsimage_grid_hwNemotronHOmniModelTestMixin._audio_features_prepare_config_and_inputsinput_features_maskFix
Add
device=torch_deviceto each of the four tensor creations. Also removes the now-redundant.to(torch_device)on thecreate_audio_maskreturn value.🤖 Generated with Claude Code