---
title: BatchSizeFinder
description: >-
  Finds the largest batch size supported by a given model before encountering an
  out of memory (OOM) error.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.callbacks.BatchSizeFinder
---
# BatchSizeFinder {#batchsizefinder}

### `class lightning.pytorch.callbacks.BatchSizeFinder(mode='power', steps_per_trial=3, init_val=2, max_trials=25, batch_arg_name='batch_size', margin=0.05, max_val=8192)` {#lightning.pytorch.callbacks.BatchSizeFinder}

Bases: [`Callback`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.Callback.md#lightning.pytorch.callbacks.Callback "lightning.pytorch.callbacks.callback.Callback")

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 [experimental](/docs/pytorch/2.6.6/reference/versioning.md#experimental-api) feature.

Parameters:

*   **mode** ([`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) –
    
    search strategy to update the batch size:
    
    *   `'power'`: Keep multiplying the batch size by 2, until we get an OOM error.
        
    *   `'binsearch'`: Initially keep multiplying by 2 and after encountering an OOM error do a binary search between the last successful batch size and the batch size that failed.
        
    
*   **steps\_per\_trial** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – number of steps to run with a given batch size. Ideally 1 should be enough to test if an OOM error occurs, however in practice a few are needed.
    
*   **init\_val** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – initial batch size to start the search with.
    
*   **max\_trials** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – max number of increases in batch size done before algorithm is terminated
    
*   **batch\_arg\_name** ([`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) –
    
    name of the attribute that stores the batch size. It is expected that the user has provided a model or datamodule that has a hyperparameter with that name. We will look for this attribute name in the following places
    
    *   `model`
        
    *   `model.hparams`
        
    *   `trainer.datamodule` (the datamodule passed to the tune method)
        
    
*   **margin** ([`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)")) – Margin to reduce the found batch size by to provide a safety buffer. Only applied when using ‘binsearch’ mode. Should be a float between 0 and 1. Defaults to 0.05 (5% reduction).
    
*   **max\_val** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Maximum batch size limit, defaults to 8192. Helps prevent testing unrealistically large or inefficient batch sizes (e.g., 2\*\*25) when running on CPU or when automatic OOM detection is not available.
    

Example:

```python
# 1. Customize the BatchSizeFinder callback to run at different epochs. This feature is
# useful while fine-tuning models since you can't always use the same batch size after
# unfreezing the backbone.
from lightning.pytorch.callbacks import BatchSizeFinder


class FineTuneBatchSizeFinder(BatchSizeFinder):
    def __init__(self, milestones, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.milestones = milestones

    def on_fit_start(self, *args, **kwargs):
        return

    def on_train_epoch_start(self, trainer, pl_module):
        if trainer.current_epoch in self.milestones or trainer.current_epoch == 0:
            self.scale_batch_size(trainer, pl_module)


trainer = Trainer(callbacks=[FineTuneBatchSizeFinder(milestones=(5, 10))])
trainer.fit(...)
```

Example:

```python
# 2. Run batch size finder for validate/test/predict.
from lightning.pytorch.callbacks import BatchSizeFinder


class EvalBatchSizeFinder(BatchSizeFinder):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def on_fit_start(self, *args, **kwargs):
        return

    def on_test_start(self, trainer, pl_module):
        self.scale_batch_size(trainer, pl_module)


trainer = Trainer(callbacks=[EvalBatchSizeFinder()])
trainer.test(...)
```

#### `on_fit_start(trainer, pl_module)` {#lightning.pytorch.callbacks.BatchSizeFinder.on_fit_start}

Called when fit begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_predict_start(trainer, pl_module)` {#lightning.pytorch.callbacks.BatchSizeFinder.on_predict_start}

Called when the predict begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_test_start(trainer, pl_module)` {#lightning.pytorch.callbacks.BatchSizeFinder.on_test_start}

Called when the test begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_validation_start(trainer, pl_module)` {#lightning.pytorch.callbacks.BatchSizeFinder.on_validation_start}

Called when the validation loop begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `setup(trainer, pl_module, stage=None)` {#lightning.pytorch.callbacks.BatchSizeFinder.setup}

Called when fit, validate, test, predict, or tune begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

---

**Navigation**

- Parent: [Api](/docs/pytorch/2.6.6/optional-api/api.md)