---
title: ModelHooks
description: Hooks to be used in LightningModule.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.core.hooks.ModelHooks
---
# ModelHooks {#modelhooks}

### `class lightning.pytorch.core.hooks.ModelHooks` {#lightning.pytorch.core.hooks.ModelHooks}

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

Hooks to be used in LightningModule.

#### `configure_model()` {#lightning.pytorch.core.hooks.ModelHooks.configure_model}

Hook to create modules in a strategy and precision aware context.

This is particularly useful for when using sharded strategies (FSDP and DeepSpeed), where we’d like to shard the model instantly to save memory and initialization time. For non-sharded strategies, you can choose to override this hook or to initialize your model under the [`init_module()`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.init_module "lightning.pytorch.trainer.trainer.Trainer.init_module") context manager.

This hook is called during each of fit/val/test/predict stages in the same process, so ensure that implementation of this hook is **idempotent**, i.e., after the first time the hook is called, subsequent calls to it should be a no-op.

Return type:

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

#### `configure_sharded_model()` {#lightning.pytorch.core.hooks.ModelHooks.configure_sharded_model}

Deprecated.

Use [`configure_model()`](#lightning.pytorch.core.hooks.ModelHooks.configure_model "lightning.pytorch.core.hooks.ModelHooks.configure_model") instead.

Return type:

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

#### `on_after_backward()` {#lightning.pytorch.core.hooks.ModelHooks.on_after_backward}

Called after `loss.backward()` and before optimizers are stepped. :rtype: [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

Note

If using native AMP, the gradients will not be unscaled at this point. Use the `on_before_optimizer_step` if you need the unscaled gradients.

#### `on_before_backward(loss)` {#lightning.pytorch.core.hooks.ModelHooks.on_before_backward}

Called before `loss.backward()`.

Parameters:

**loss** ([`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")) – Loss divided by number of batches for gradient accumulation and scaled if using AMP.

Return type:

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

#### `on_before_optimizer_step(optimizer)` {#lightning.pytorch.core.hooks.ModelHooks.on_before_optimizer_step}

Called before `optimizer.step()`.

If using gradient accumulation, the hook is called once the gradients have been accumulated. See: [`accumulate_grad_batches`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.params.accumulate_grad_batches "lightning.pytorch.trainer.trainer.Trainer").

If using AMP, the loss will be unscaled before calling this hook. See these [docs](https://pytorch.org/docs/stable/notes/amp_examples.html#working-with-unscaled-gradients) for more information on the scaling of gradients.

If clipping gradients, the gradients will not have been clipped yet.

Parameters:

**optimizer** ([`Optimizer`](https://docs.pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.15)")) – Current optimizer being used.

Return type:

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

Example:

```python
def on_before_optimizer_step(self, optimizer):
    # example to inspect gradient information in tensorboard
    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
        for k, v in self.named_parameters():
            self.logger.experiment.add_histogram(
                tag=k, values=v.grad, global_step=self.trainer.global_step
            )
```

#### `on_before_zero_grad(optimizer)` {#lightning.pytorch.core.hooks.ModelHooks.on_before_zero_grad}

Called after `training_step()` and before `optimizer.zero_grad()`.

Called in the training loop after taking an optimizer step and before zeroing grads. Good place to inspect weight information with weights updated.

This is where it is called:

```python
for optimizer in optimizers:
    out = training_step(...)

    model.on_before_zero_grad(optimizer) # < ---- called here
    optimizer.zero_grad()

    backward()
```

Parameters:

**optimizer** ([`Optimizer`](https://docs.pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.15)")) – The optimizer for which grads should be zeroed.

Return type:

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

#### `on_fit_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_fit_end}

Called at the very end of fit.

If on DDP it is called on every process

Return type:

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

#### `on_fit_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_fit_start}

Called at the very beginning of fit.

If on DDP it is called on every process

Return type:

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

#### `on_predict_batch_end(outputs, batch, batch_idx, dataloader_idx=0)` {#lightning.pytorch.core.hooks.ModelHooks.on_predict_batch_end}

Called in the predict loop after the batch.

Parameters:

*   **outputs** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – The outputs of predict\_step(x)
    
*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – The batched data as it is returned by the prediction DataLoader.
    
*   **batch\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the batch
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the dataloader
    

Return type:

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

#### `on_predict_batch_start(batch, batch_idx, dataloader_idx=0)` {#lightning.pytorch.core.hooks.ModelHooks.on_predict_batch_start}

Called in the predict loop before anything happens for that batch.

Parameters:

*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – The batched data as it is returned by the test DataLoader.
    
*   **batch\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the batch
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the dataloader
    

Return type:

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

#### `on_predict_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_predict_end}

Called at the end of predicting.

Return type:

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

#### `on_predict_epoch_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_predict_epoch_end}

Called at the end of predicting.

Return type:

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

#### `on_predict_epoch_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_predict_epoch_start}

Called at the beginning of predicting.

Return type:

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

#### `on_predict_model_eval()` {#lightning.pytorch.core.hooks.ModelHooks.on_predict_model_eval}

Called when the predict loop starts.

The predict loop by default calls `.eval()` on the LightningModule before it starts. Override this hook to change the behavior.

Return type:

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

#### `on_predict_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_predict_start}

Called at the beginning of predicting.

Return type:

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

#### `on_test_batch_end(outputs, batch, batch_idx, dataloader_idx=0)` {#lightning.pytorch.core.hooks.ModelHooks.on_test_batch_end}

Called in the test loop after the batch.

Parameters:

*   **outputs** (`Union`\[[`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)"), [`Mapping`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Mapping "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – The outputs of test\_step(x)
    
*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – The batched data as it is returned by the test DataLoader.
    
*   **batch\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the batch
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the dataloader
    

Return type:

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

#### `on_test_batch_start(batch, batch_idx, dataloader_idx=0)` {#lightning.pytorch.core.hooks.ModelHooks.on_test_batch_start}

Called in the test loop before anything happens for that batch.

Parameters:

*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – The batched data as it is returned by the test DataLoader.
    
*   **batch\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the batch
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the dataloader
    

Return type:

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

#### `on_test_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_test_end}

Called at the end of testing.

Return type:

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

#### `on_test_epoch_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_test_epoch_end}

Called in the test loop at the very end of the epoch.

Return type:

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

#### `on_test_epoch_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_test_epoch_start}

Called in the test loop at the very beginning of the epoch.

Return type:

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

#### `on_test_model_eval()` {#lightning.pytorch.core.hooks.ModelHooks.on_test_model_eval}

Called when the test loop starts.

The test loop by default calls `.eval()` on the LightningModule before it starts. Override this hook to change the behavior. See also [`on_test_model_train()`](#lightning.pytorch.core.hooks.ModelHooks.on_test_model_train "lightning.pytorch.core.hooks.ModelHooks.on_test_model_train").

Return type:

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

#### `on_test_model_train()` {#lightning.pytorch.core.hooks.ModelHooks.on_test_model_train}

Called when the test loop ends.

The test loop by default restores the training mode of the LightningModule to what it was before starting testing. Override this hook to change the behavior. See also [`on_test_model_eval()`](#lightning.pytorch.core.hooks.ModelHooks.on_test_model_eval "lightning.pytorch.core.hooks.ModelHooks.on_test_model_eval").

Return type:

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

#### `on_test_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_test_start}

Called at the beginning of testing.

Return type:

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

#### `on_train_batch_end(outputs, batch, batch_idx)` {#lightning.pytorch.core.hooks.ModelHooks.on_train_batch_end}

Called in the training loop after the batch.

Parameters:

*   **outputs** (`Union`\[[`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)"), [`Mapping`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Mapping "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – The outputs of training\_step(x)
    
*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – The batched data as it is returned by the training DataLoader.
    
*   **batch\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the batch
    

Return type:

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

Note

The value `outputs["loss"]` here will be the normalized value w.r.t `accumulate_grad_batches` of the loss returned from `training_step`.

#### `on_train_batch_start(batch, batch_idx)` {#lightning.pytorch.core.hooks.ModelHooks.on_train_batch_start}

Called in the training loop before anything happens for that batch.

If you return -1 here, you will skip training for the rest of the current epoch. Learning rate scheduler will still be stepped at the end of epoch.

Parameters:

*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – The batched data as it is returned by the training DataLoader.
    
*   **batch\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the batch
    

Return type:

[`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\]

#### `on_train_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_train_end}

Called at the end of training before logger experiment is closed.

Return type:

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

#### `on_train_epoch_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_train_epoch_end}

Called in the training loop at the very end of the epoch.

To access all batch outputs at the end of the epoch, you can cache step outputs as an attribute of the `LightningModule` and access them in this hook:

```python
class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()
```

Return type:

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

#### `on_train_epoch_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_train_epoch_start}

Called in the training loop at the very beginning of the epoch.

Return type:

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

#### `on_train_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_train_start}

Called at the beginning of training after sanity check.

Return type:

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

#### `on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx=0)` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_batch_end}

Called in the validation loop after the batch.

Parameters:

*   **outputs** (`Union`\[[`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)"), [`Mapping`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Mapping "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – The outputs of validation\_step(x)
    
*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – The batched data as it is returned by the validation DataLoader.
    
*   **batch\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the batch
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the dataloader
    

Return type:

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

#### `on_validation_batch_start(batch, batch_idx, dataloader_idx=0)` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_batch_start}

Called in the validation loop before anything happens for that batch.

Parameters:

*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – The batched data as it is returned by the validation DataLoader.
    
*   **batch\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the batch
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – the index of the dataloader
    

Return type:

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

#### `on_validation_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_end}

Called at the end of validation.

Return type:

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

#### `on_validation_epoch_end()` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_epoch_end}

Called in the validation loop at the very end of the epoch.

Return type:

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

#### `on_validation_epoch_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_epoch_start}

Called in the validation loop at the very beginning of the epoch.

Return type:

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

#### `on_validation_model_eval()` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_model_eval}

Called when the validation loop starts.

The validation loop by default calls `.eval()` on the LightningModule before it starts. Override this hook to change the behavior. See also [`on_validation_model_train()`](#lightning.pytorch.core.hooks.ModelHooks.on_validation_model_train "lightning.pytorch.core.hooks.ModelHooks.on_validation_model_train").

Return type:

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

#### `on_validation_model_train()` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_model_train}

Called when the validation loop ends.

The validation loop by default restores the training mode of the LightningModule to what it was before starting validation. Override this hook to change the behavior. See also [`on_validation_model_eval()`](#lightning.pytorch.core.hooks.ModelHooks.on_validation_model_eval "lightning.pytorch.core.hooks.ModelHooks.on_validation_model_eval").

Return type:

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

#### `on_validation_model_zero_grad()` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_model_zero_grad}

Called by the training loop to release gradients before entering the validation loop.

Return type:

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

#### `on_validation_start()` {#lightning.pytorch.core.hooks.ModelHooks.on_validation_start}

Called at the beginning of validation.

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)