mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Chenjie Luo
3ad4f4f093
## Summary - `get_max_batch_size` halved `target_data_batch` on `torch.cuda.OutOfMemoryError` but never rebuilt `target_input`, so each retry re-fed the same too-large tensor — the retry loop was effectively a no-op. - Refactor the expand logic into an `_expand_to(batch)` helper, rebuild `target_input` after halving, and call `torch.cuda.empty_cache()` between attempts. ## Test plan - [x] New unit test `test_get_max_batch_size_oom_retry_shrinks_input` mocks `torch.cuda.*` and asserts the second retry receives the halved tensor (shapes seen: `[1, 10, 5]`, regulated result `4`). - [x] `pytest tests/unit/torch/utils/test_dataset_utils.py` — 14/14 pass (skipping the network-only minipile test). - [x] `pre-commit` (ruff, mypy, bandit, license headers) clean on commit. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Enhanced GPU memory management during batch size detection. When out-of-memory errors occur during the initial probing phase, the system now properly adapts input tensors to smaller batch sizes and clears GPU cache before retry attempts, resulting in more reliable recovery and stable batch sizing across diverse hardware environments. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>