You may consider using Composer by MosaicML.

It has features like automatically calculating batch size to prevent CUDA out of memory errors and early stopping.

The library uses a simple Trainer with customizable callbacks to modify any part of the training loop.