Feature to automatically choose batch size
Let's add a flag:
# default False
Trainer(auto_scale_batch_size=True)
This should do  binary search on batch size:

Run a few train steps using the current batch size.
if OOM batch_size / 2.
If no OOM, batch_size = batch_size * 1.5.

And so on until we find the optimal batch size. At this point log it so the user knows (including tensorboard), and continue training with the new batch size.
Ideally the user fixes the batch size in future runs to tune the learning rate.