[Fix] Re-expand target_input on OOM in get_max_batch_size (NVIDIA#1374)

## 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>

Author: cjluo-nv

modelopt/torch/utils/dataset_utils.py

@@ -519,12 +519,16 @@ def _get_free_gpu_mem():
             target_data_batch = 1  # pragma: no cover
         else:
             target_data_batch = max(int(free_mem_before / mem_diff_per_data_batch), 1)
-        target_input = sample_input_single_batch.expand(
-            [
-                target_data_batch if index == 0 else dim
-                for index, dim in enumerate(sample_input_single_batch.shape)
-            ]
-        )
+
+        def _expand_to(batch: int) -> torch.Tensor:
+            return sample_input_single_batch.expand(
+                [
+                    batch if index == 0 else dim
+                    for index, dim in enumerate(sample_input_single_batch.shape)
+                ]
+            )
+
+        target_input = _expand_to(target_data_batch)
 
         # For some models on multi GPU, we observe the memory per batch is not a constant.
         # So we just test the target batch size and make sure we do not go OOM.
@@ -535,6 +539,8 @@ def _get_free_gpu_mem():
                     break
                 except torch.cuda.OutOfMemoryError:  # pragma: no cover - GPU OOM retry path
                     target_data_batch = target_data_batch // 2  # pragma: no cover
+                    target_input = _expand_to(target_data_batch)  # pragma: no cover
+                    torch.cuda.empty_cache()  # pragma: no cover
 
     # Regulate the data batch target to be 1, 2, 4, 8, 12, ..., capped at 64
     if target_data_batch < 2:

tests/unit/torch/utils/test_dataset_utils.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 import torch
@@ -24,6 +24,7 @@
     _forward_loop,
     _process_batch,
     get_dataset_samples,
+    get_max_batch_size,
 )
 
 
@@ -231,6 +232,54 @@ def test_disable_use_cache_restores_on_exception():
     assert model.config.use_cache is True
 
 
+def test_get_max_batch_size_oom_retry_shrinks_input():
+    """On OOM, target_input must be re-expanded to the halved batch size.
+
+    Regression test: previously target_input was built once and never shrunk,
+    so OOM retries kept feeding the same too-large tensor to infer_method.
+    """
+    seq_len = 8
+    sample_input = torch.ones((1, seq_len), dtype=torch.int32)
+
+    seen_batch_sizes: list[int] = []
+
+    def fake_forward(x):
+        seen_batch_sizes.append(x.shape[0])
+        # First call is the single-batch probe — succeeds.
+        # Second call is the target-batch attempt — OOMs.
+        # Third call (after halving) — succeeds.
+        if len(seen_batch_sizes) == 2:
+            raise torch.cuda.OutOfMemoryError
+
+    model = Mock(spec=torch.nn.Module)
+    model.forward = fake_forward
+    model.__class__.__name__ = "DummyModel"  # not enc/dec
+
+    free_before = 1000
+    free_after = 900  # 100 bytes per batch -> target = 1000/100 = 10
+
+    device_props = Mock()
+    device_props.total_memory = 10**12
+
+    with (
+        patch("torch.cuda.empty_cache"),
+        patch("torch.cuda.get_device_properties", return_value=device_props),
+        patch("torch.cuda.device_count", return_value=1),
+        patch("torch.cuda.mem_get_info", side_effect=[(free_before, 0), (free_after, 0)]),
+        patch("torch.cuda.max_memory_allocated", side_effect=[0, 0]),
+    ):
+        result = get_max_batch_size(
+            model,
+            max_sample_length=seq_len,
+            sample_input_single_batch=sample_input,
+        )
+
+    # Forward calls: probe(1), retry-at-target(10), retry-after-halve(5)
+    assert seen_batch_sizes == [1, 10, 5]
+    # Final batch is 5 -> regulated to 4 (5 // 4 * 4 = 4).
+    assert result == 4
+
+
 @pytest.mark.parametrize("test_local_path", [True, False])
 def test_get_dataset_samples_with_unsupported_minipile_dataset(tmp_path, test_local_path):
     pytest.importorskip("datasets")