diff --git a/src/lightning/pytorch/callbacks/batch_size_finder.py b/src/lightning/pytorch/callbacks/batch_size_finder.py index c95803c035db0..e348f4946e51b 100644 --- a/src/lightning/pytorch/callbacks/batch_size_finder.py +++ b/src/lightning/pytorch/callbacks/batch_size_finder.py @@ -31,10 +31,11 @@ class BatchSizeFinder(Callback): - """The ``BatchSizeFinder`` callback tries to find the largest batch size for a given model that does not give an - out of memory (OOM) error. All you need to do is add it as a callback inside Trainer and call - ``trainer.{fit,validate,test,predict}``. Internally it calls the respective step function ``steps_per_trial`` times - for each batch size until one of the batch sizes generates an OOM error. + """Finds the largest batch size supported by a given model before encountering an out of memory (OOM) error. + + All you need to do is add it as a callback inside Trainer and call ``trainer.{fit,validate,test,predict}``. + Internally, it calls the respective step function ``steps_per_trial`` times for each batch size until one + of the batch sizes generates an OOM error. .. warning:: This is an :ref:`experimental ` feature.