Lightning-AI/pytorch-lightning · error · NotImplementedError

The `PredictionWriterCallback` does not support using `datal

Error message

The `PredictionWriterCallback` does not support using `dataloader_iter`.

What it means

BasePredictionWriter streams predictions batch-by-batch, which requires the batch itself as an argument to write_on_batch_end. If the LightningModule's predict_step uses the dataloader_iter style signature (iterating the dataloater manually), the batch is not available to the writer, so setup() raises NotImplementedError before prediction starts.

Source

Thrown at src/lightning/pytorch/callbacks/prediction_writer.py:118

        # 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)

    """

    def __init__(self, write_interval: Literal["batch", "epoch", "batch_and_epoch"] = "batch") -> None:
        if write_interval not in list(WriteInterval):
            raise MisconfigurationException(f"`write_interval` should be one of {[i.value for i in WriteInterval]}.")
        self.interval = WriteInterval(write_interval)

    @override
    def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
        if is_param_in_hook_signature(pl_module.predict_step, "dataloader_iter", explicit=True):
            raise NotImplementedError("The `PredictionWriterCallback` does not support using `dataloader_iter`.")

    def write_on_batch_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        prediction: Any,
        batch_indices: Optional[Sequence[int]],
        batch: Any,
        batch_idx: int,
        dataloader_idx: int,
    ) -> None:
        """Override with the logic to write a single batch."""
        raise NotImplementedError()

    def write_on_epoch_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",

View on GitHub (pinned to 9fed5c27d2)

Solutions

Refactor predict_step to accept the batch directly (standard style) so the writer receives predictions per batch
If dataloader_iter is required, write predictions manually inside predict_step (e.g., to disk or a buffer) instead of using BasePredictionWriter
Subclass the writer and override write_on_batch_end only if you can guarantee batch access another way

Example fix

# before
class M(pl.LightningModule):
    def predict_step(self, dataloader_iter):
        batch, _ = next(dataloader_iter)
        return self(batch)
# after
class M(pl.LightningModule):
    def predict_step(self, batch, batch_idx):
        return self(batch)

Defensive patterns

Strategy: type-guard

Validate before calling

import inspect
sig = inspect.signature(model.predict_step)
assert 'dataloader_iter' not in sig.parameters, 'PredictionWriter requires standard predict_step(batch, batch_idx)'

Type guard

def supports_prediction_writer(model) -> bool:
    return 'dataloader_iter' not in inspect.signature(model.predict_step).parameters

Try / catch

try:
    writer.setup(trainer, model, 'predict')
except NotImplementedError:
    writer = None  # handle predictions manually in predict_step

Prevention

Avoid dataloader_iter style when using prediction callbacks
Keep predict_step signatures standard (batch, batch_idx)

When it happens

Trigger: Defining predict_step(self, dataloader_iter, ...) on your LightningModule (or a LightningDataModule providing that style) and then attaching a BasePredictionWriter subclass; the check uses is_param_in_hook_signature with explicit=True, so even an explicitly named dataloader_iter parameter triggers it.

Common situations: Using dataloader_iter to access batch indices or do custom batching during prediction, then adding a prediction writer callback; inheriting a shared model base class whose predict_step takes dataloader_iter.

Related errors

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/fa5b3df1be303f35. Report an issue: GitHub.