Improve recovery from OOM errors mechanism

**Is your feature request related to a problem? Please describe.**

Inference and training have a mechanism to catch OOM errors. When this happens, ideally all dangling GPU memory is freed up, then parameters determining the peak memory usage are decreased, and the task is repeated.

The current approach:

* Is incomplete: Not all GPU memory is released
* Is limited to certain parameters, which are present with naive SDPA only

**Describe the solution you'd like**

Should figure out how to release all GPU memory allocated inside the block to be repeated. Also, support reducing other parameters, in particular the batch size.

Start by locating the existing OOM handling in the inference and training paths, then trace how the repeated block releases GPU memory and selects parameters. Done means the mechanism releases all memory allocated inside the retried block and can reduce additional parameters, including batch size, before repeating the task.

Written by the indexing model from the issue text.