---
title: BasePredictionWriter
description: Base class to implement how the predictions should be stored.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.callbacks.BasePredictionWriter
---
# BasePredictionWriter {#basepredictionwriter}

### `class lightning.pytorch.callbacks.BasePredictionWriter(write_interval='batch')` {#lightning.pytorch.callbacks.BasePredictionWriter}

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

Base class to implement how the predictions should be stored.

Parameters:

**write\_interval** ([`Literal`](https://docs.python.org/3/library/typing.html#typing.Literal "(in Python v3.15)")\[`'batch'`, `'epoch'`, `'batch_and_epoch'`\]) – When to write.

Example:

```python
import torch
from lightning.pytorch.callbacks import BasePredictionWriter

class CustomWriter(BasePredictionWriter):

    def __init__(self, output_dir, write_interval):
        super().__init__(write_interval)
        self.output_dir = output_dir

    def write_on_batch_end(
        self, trainer, pl_module, prediction, batch_indices, batch, batch_idx, dataloader_idx
    ):
        torch.save(prediction, os.path.join(self.output_dir, dataloader_idx, f"{batch_idx}.pt"))

    def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
        torch.save(predictions, os.path.join(self.output_dir, "predictions.pt"))


pred_writer = CustomWriter(output_dir="pred_path", write_interval="epoch")
trainer = Trainer(callbacks=[pred_writer])
model = BoringModel()
trainer.predict(model, return_predictions=False)
```

Example:

```python
# multi-device inference example

import torch
from lightning.pytorch.callbacks import BasePredictionWriter

class CustomWriter(BasePredictionWriter):

    def __init__(self, output_dir, write_interval):
        super().__init__(write_interval)
        self.output_dir = output_dir

    def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
        # this will create N (num processes) files in `output_dir` each containing
        # the predictions of it's respective rank
        torch.save(predictions, os.path.join(self.output_dir, f"predictions_{trainer.global_rank}.pt"))

        # optionally, you can also save `batch_indices` to get the information about the data index
        # from your prediction data
        torch.save(batch_indices, os.path.join(self.output_dir, f"batch_indices_{trainer.global_rank}.pt"))


# or you can set `write_interval="batch"` and override `write_on_batch_end` to save
# predictions at batch level
pred_writer = CustomWriter(output_dir="pred_path", write_interval="epoch")
trainer = Trainer(accelerator="gpu", strategy="ddp", devices=8, callbacks=[pred_writer])
model = BoringModel()
trainer.predict(model, return_predictions=False)
```

#### `on_predict_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)` {#lightning.pytorch.callbacks.BasePredictionWriter.on_predict_batch_end}

Called when the predict batch ends.

Return type:

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

#### `on_predict_epoch_end(trainer, pl_module)` {#lightning.pytorch.callbacks.BasePredictionWriter.on_predict_epoch_end}

Called when the predict epoch ends.

Return type:

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

#### `setup(trainer, pl_module, stage)` {#lightning.pytorch.callbacks.BasePredictionWriter.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)")

#### `write_on_batch_end(trainer, pl_module, prediction, batch_indices, batch, batch_idx, dataloader_idx)` {#lightning.pytorch.callbacks.BasePredictionWriter.write_on_batch_end}

Override with the logic to write a single batch.

Return type:

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

#### `write_on_epoch_end(trainer, pl_module, predictions, batch_indices)` {#lightning.pytorch.callbacks.BasePredictionWriter.write_on_epoch_end}

Override with the logic to write all batches.

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)