Free cached memory inside the batch-size probe, not after estimation

[avnoordenne]

determine_max_batch_size deliberately grows batches until the GPU runs out of memory, but left PyTorch's caching allocator holding the device on the way out. #589 added an empty_cache() for this, at the end of estimate_max_memory_scaler — but InFlightAutoBatcher._get_first_batch calls _get_next_states() between the first probe and estimate_max_memory_scaler, so the starved allocation happens one step before the guard:

1. determine_max_batch_size(first_state) — probe, leaves the card full
2. _get_next_states() — neighbor-list calls ← fails here
3. estimate_max_memory_scaler() — two more probes, then fix(autobatching): recognise Warp OOM in memory estimation and free cache before the real run #589's empty_cache()
4. _get_next_states() again — protected
5. the optimization run — protected

This is only reachable with memory_scales_with="n_edges" and alchemiops installed. n_atoms and the default n_atoms_x_density compute their scalers with pure tensor arithmetic and never build a neighbor list, so on the default path _get_next_states() makes no Warp allocations at all.

Move the release into determine_max_batch_size itself, in a finally covering all three exits: normal completion, the early return on a caught OOM, and the re-raise on a non-OOM error. The OOM return is the path that matters in practice. Deferring to the outer finally matters: by then the except clause's implicit del exc has dropped the traceback frames and the activations they pin, so empty_cache() can actually reclaim them.

Measured with SevenNet 7net-0 over 878 molecules on a 24 GB card, memory_scales_with="n_edges". The ladder is seeded from a 16-atom molecule and capped at the set's 22,260 atoms, so it tops out at 1071 replicas (17,136 atoms), which OOMs:

after the probe returns	before	after
reserved	21.27 GB	0.05 GB
free on device	0.03 GB	21.24 GB
next neighbor-list call	RuntimeError: Failed to allocate 72 bytes on device 'cuda:0'	OK
returned batch size	350	350

The estimate is unchanged, and can't change: measure_model_memory_forward already calls empty_cache() at the top of every rung, so each measurement already begins from a cleared cache, and the value it returns is max_memory_allocated() — a high-water mark of allocated bytes, which empty_cache() cannot move.

The guard in estimate_max_memory_scaler is now redundant and is removed, so the invariant lives with the function that breaks it.

Summary

• Release PyTorch's caching allocator inside determine_max_batch_size via finally, so every exit path hands the device back — fixes Failed to allocate N bytes from Warp-backed neighbor lists in _get_next_states() when memory_scales_with="n_edges"
• Remove the now-redundant empty_cache() from estimate_max_memory_scaler, whose probes now clean up after themselves
• Add a parametrized regression test covering both the OOM and clean-completion exits