{"record":{"id":"3ec6e68b3429b6a0","repo":"Lightning-AI/pytorch-lightning","slug":"couldn-t-infer-the-batch-indices-fetched-from-your","errorCode":null,"errorMessage":"Couldn't infer the batch indices fetched from your dataloader: `{type(dataloader).__name__}`","messagePattern":"Couldn't infer the batch indices fetched from your dataloader: `(.+?)`","errorType":"console","errorClass":null,"httpStatus":null,"severity":"warning","filePath":"src/lightning/pytorch/loops/prediction_loop.py","lineNumber":307,"sourceCode":"        step_kwargs = OrderedDict([(\"batch\", batch), (\"batch_idx\", batch_idx)])\n        if dataloader_idx is not None:\n            step_kwargs[\"dataloader_idx\"] = dataloader_idx\n        return step_kwargs\n\n    def _build_step_args_from_hook_kwargs(self, hook_kwargs: OrderedDict, step_hook_name: str) -> tuple:\n        \"\"\"Helper method to build args for `predict_step`.\"\"\"\n        kwargs = hook_kwargs.copy()\n        step_hook_fx = getattr(self.trainer.lightning_module, step_hook_name)\n        if not is_param_in_hook_signature(step_hook_fx, \"batch_idx\", min_args=2):\n            kwargs.pop(\"batch_idx\", None)\n        return tuple(kwargs.values())\n\n    def _get_batch_indices(self, dataloader: object) -> list[list[int]]:  # batches x samples\n        \"\"\"Returns a reference to the seen batch indices if the dataloader has a batch sampler wrapped by our\n        :class:`~lightning.pytorch.overrides.distributed._IndexBatchSamplerWrapper`.\"\"\"\n        batch_sampler = getattr(dataloader, \"batch_sampler\", None)\n        if not isinstance(batch_sampler, _IndexBatchSamplerWrapper):\n            self._warning_cache.warn(\n                f\"Couldn't infer the batch indices fetched from your dataloader: `{type(dataloader).__name__}`\"\n            )\n            return []\n        return batch_sampler.seen_batch_indices\n\n    def _store_data_for_prediction_writer(self, batch_idx: int, dataloader_idx: int) -> bool:\n        prediction_writers = [cb for cb in self.trainer.callbacks if isinstance(cb, BasePredictionWriter)]\n        any_on_epoch = any(cb.interval.on_epoch for cb in prediction_writers)\n        any_on_batch = any(cb.interval.on_batch for cb in prediction_writers)\n        if any_on_batch or any_on_epoch:\n            combined_loader = self._combined_loader\n            assert combined_loader is not None\n            dataloader = combined_loader.flattened[dataloader_idx]\n            batch_indices = self._get_batch_indices(dataloader)\n            if not batch_indices:\n                # this is only available with `_IndexBatchSamplerWrapper`, but it's only used on DataLoaders, if this is\n                # reached, it's likely because a non-DataLoader was passed\n                return any_on_epoch","sourceCodeStart":289,"sourceCodeEnd":325,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/loops/prediction_loop.py#L289-L325","documentation":"_get_batch_indices needs the dataloader's batch_sampler to be Lightning's _IndexBatchSamplerWrapper to know which sample indices each batch contained. With a plain/unwrapped dataloader it warns and returns [], which disables per-sample index tracking for the prediction writer.","triggerScenarios":"Passing a custom/collated dataloader (e.g. a DataLoader whose batch_sampler is a standard sampler) to trainer.predict with a callback relying on batch indices; manually constructed dataloaders bypassing Lightning's wrapping.","commonSituations":"Custom IterableDataset or a DataLoader built outside trainer.fit/predict setup hooks so Lightning never wrapped its sampler.","solutions":["Let Lightning create the dataloader via predict_dataloader/on_predict_dataloader hooks or pass the DataLoader object to trainer.predict so it gets wrapped","If using a plain iterator, don't rely on batch indices in your prediction writer","Track indices yourself by yielding (x, idx) from the dataset"],"exampleFix":"# before\ntrainer.predict(model, dataloaders=DataLoader(ds, batch_sampler=my_sampler()))  # unwrapped\n# after\nclass DM(LightningDataModule):\n    def predict_dataloader(self):\n        return DataLoader(ds, batch_size=4)  # Lightning wraps the sampler","handlingStrategy":"fallback","validationCode":"from lightning.pytorch.overrides.distributed import _IndexBatchSamplerWrapper\nassert isinstance(getattr(dl, 'batch_sampler', None), _IndexBatchSamplerWrapper), 'sampler not wrapped; indices unavailable'","typeGuard":"def has_wrapped_sampler(dl) -> bool:\n    from lightning.pytorch.overrides.distributed import _IndexBatchSamplerWrapper\n    return isinstance(getattr(dl, 'batch_sampler', None), _IndexBatchSamplerWrapper)","tryCatchPattern":null,"preventionTips":["Provide dataloaders via Lightning hooks so they get wrapped","Carry indices in the dataset if you need per-sample tracking"],"tags":["prediction","dataloader","batch-indices","lightning"],"backgroundTag":"unwrapped-batch-sampler","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}