fix early stop

[Tetracarbonylnickel]

Previously, using break only exited the training loop but did not close the surrounding with context. As a result, early stopping did not terminate training properly.
This change ensures that early stopping exits the context manager as well, allowing training to terminate cleanly.

[PythonFZ]

Seems like a good change. There are also some other restructurings in the training loop that look good, but I'd suggest either now or in a future PR to refactor this a little bit and see if e.g. train / val steps can be written as separate functions making changes easier to review.