Lightning-AI/pytorch-lightning · warning

Couldn't infer the batch indices fetched from your dataloade

Error message

Couldn't infer the batch indices fetched from your dataloader: `{type(dataloader).__name__}`

What it means

_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.

Source

Thrown at src/lightning/pytorch/loops/prediction_loop.py:307

        step_kwargs = OrderedDict([("batch", batch), ("batch_idx", batch_idx)])
        if dataloader_idx is not None:
            step_kwargs["dataloader_idx"] = dataloader_idx
        return step_kwargs

    def _build_step_args_from_hook_kwargs(self, hook_kwargs: OrderedDict, step_hook_name: str) -> tuple:
        """Helper method to build args for `predict_step`."""
        kwargs = hook_kwargs.copy()
        step_hook_fx = getattr(self.trainer.lightning_module, step_hook_name)
        if not is_param_in_hook_signature(step_hook_fx, "batch_idx", min_args=2):
            kwargs.pop("batch_idx", None)
        return tuple(kwargs.values())

    def _get_batch_indices(self, dataloader: object) -> list[list[int]]:  # batches x samples
        """Returns a reference to the seen batch indices if the dataloader has a batch sampler wrapped by our
        :class:`~lightning.pytorch.overrides.distributed._IndexBatchSamplerWrapper`."""
        batch_sampler = getattr(dataloader, "batch_sampler", None)
        if not isinstance(batch_sampler, _IndexBatchSamplerWrapper):
            self._warning_cache.warn(
                f"Couldn't infer the batch indices fetched from your dataloader: `{type(dataloader).__name__}`"
            )
            return []
        return batch_sampler.seen_batch_indices

    def _store_data_for_prediction_writer(self, batch_idx: int, dataloader_idx: int) -> bool:
        prediction_writers = [cb for cb in self.trainer.callbacks if isinstance(cb, BasePredictionWriter)]
        any_on_epoch = any(cb.interval.on_epoch for cb in prediction_writers)
        any_on_batch = any(cb.interval.on_batch for cb in prediction_writers)
        if any_on_batch or any_on_epoch:
            combined_loader = self._combined_loader
            assert combined_loader is not None
            dataloader = combined_loader.flattened[dataloader_idx]
            batch_indices = self._get_batch_indices(dataloader)
            if not batch_indices:
                # this is only available with `_IndexBatchSamplerWrapper`, but it's only used on DataLoaders, if this is
                # reached, it's likely because a non-DataLoader was passed
                return any_on_epoch

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Let Lightning create the dataloader via predict_dataloader/on_predict_dataloader hooks or pass the DataLoader object to trainer.predict so it gets wrapped
2. If using a plain iterator, don't rely on batch indices in your prediction writer
3. Track indices yourself by yielding (x, idx) from the dataset

Example fix

# before
trainer.predict(model, dataloaders=DataLoader(ds, batch_sampler=my_sampler()))  # unwrapped
# after
class DM(LightningDataModule):
    def predict_dataloader(self):
        return DataLoader(ds, batch_size=4)  # Lightning wraps the sampler

from lightning.pytorch.overrides.distributed import _IndexBatchSamplerWrapper
assert isinstance(getattr(dl, 'batch_sampler', None), _IndexBatchSamplerWrapper), 'sampler not wrapped; indices unavailable'

def has_wrapped_sampler(dl) -> bool:
    from lightning.pytorch.overrides.distributed import _IndexBatchSamplerWrapper
    return isinstance(getattr(dl, 'batch_sampler', None), _IndexBatchSamplerWrapper)

When it happens

Trigger: 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.

Common situations: Custom IterableDataset or a DataLoader built outside trainer.fit/predict setup hooks so Lightning never wrapped its sampler.

Related errors

• You seem to have configured a sampler in your DataLoader whi
• f"An invalid dataloader was passed to `Trainer.{trainer_fn.v
• f"An invalid dataloader was passed to `Trainer.{trainer_fn.v
• train_dataloader yielded None. If this was on purpose, ignor
• Received multiple values for {', '.join(duplicated_plugin_ke

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