{"record":{"id":"fa5b3df1be303f35","repo":"Lightning-AI/pytorch-lightning","slug":"the-predictionwritercallback-does-not-support-us","errorCode":null,"errorMessage":"The `PredictionWriterCallback` does not support using `dataloader_iter`.","messagePattern":"The `PredictionWriterCallback` does not support using `dataloader_iter`\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"src/lightning/pytorch/callbacks/prediction_writer.py","lineNumber":118,"sourceCode":"\n        # or you can set `write_interval=\"batch\"` and override `write_on_batch_end` to save\n        # predictions at batch level\n        pred_writer = CustomWriter(output_dir=\"pred_path\", write_interval=\"epoch\")\n        trainer = Trainer(accelerator=\"gpu\", strategy=\"ddp\", devices=8, callbacks=[pred_writer])\n        model = BoringModel()\n        trainer.predict(model, return_predictions=False)\n\n    \"\"\"\n\n    def __init__(self, write_interval: Literal[\"batch\", \"epoch\", \"batch_and_epoch\"] = \"batch\") -> None:\n        if write_interval not in list(WriteInterval):\n            raise MisconfigurationException(f\"`write_interval` should be one of {[i.value for i in WriteInterval]}.\")\n        self.interval = WriteInterval(write_interval)\n\n    @override\n    def setup(self, trainer: \"pl.Trainer\", pl_module: \"pl.LightningModule\", stage: str) -> None:\n        if is_param_in_hook_signature(pl_module.predict_step, \"dataloader_iter\", explicit=True):\n            raise NotImplementedError(\"The `PredictionWriterCallback` does not support using `dataloader_iter`.\")\n\n    def write_on_batch_end(\n        self,\n        trainer: \"pl.Trainer\",\n        pl_module: \"pl.LightningModule\",\n        prediction: Any,\n        batch_indices: Optional[Sequence[int]],\n        batch: Any,\n        batch_idx: int,\n        dataloader_idx: int,\n    ) -> None:\n        \"\"\"Override with the logic to write a single batch.\"\"\"\n        raise NotImplementedError()\n\n    def write_on_epoch_end(\n        self,\n        trainer: \"pl.Trainer\",\n        pl_module: \"pl.LightningModule\",","sourceCodeStart":100,"sourceCodeEnd":136,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/callbacks/prediction_writer.py#L100-L136","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nclass M(pl.LightningModule):\n    def predict_step(self, dataloader_iter):\n        batch, _ = next(dataloader_iter)\n        return self(batch)\n# after\nclass M(pl.LightningModule):\n    def predict_step(self, batch, batch_idx):\n        return self(batch)","handlingStrategy":"type-guard","validationCode":"import inspect\nsig = inspect.signature(model.predict_step)\nassert 'dataloader_iter' not in sig.parameters, 'PredictionWriter requires standard predict_step(batch, batch_idx)'","typeGuard":"def supports_prediction_writer(model) -> bool:\n    return 'dataloader_iter' not in inspect.signature(model.predict_step).parameters","tryCatchPattern":"try:\n    writer.setup(trainer, model, 'predict')\nexcept NotImplementedError:\n    writer = None  # handle predictions manually in predict_step","preventionTips":["Avoid dataloader_iter style when using prediction callbacks","Keep predict_step signatures standard (batch, batch_idx)"],"tags":["pytorch-lightning","prediction-writer","dataloader-iter","unsupported-operation"],"backgroundTag":"unsupported-api-combination","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}