Lightning-AI/pytorch-lightning · error · TypeError

Lightning can't inject a (distributed) sampler into your ba

Error message

 Lightning can't inject a (distributed) sampler into your batch sampler, because it doesn't subclass PyTorch's `BatchSampler`. To mitigate this, either follow the API of `BatchSampler` or set `Trainer(use_distributed_sampler=False)`. If you choose the latter, you will be responsible for handling the distributed sampling within your batch sampler.

What it means

Raised as a TypeError when the DataLoader's batch_sampler object is not a PyTorch BatchSampler at all, so Lightning has no mechanism to inject a distributed sampler or adjust drop_last. This is the 'we don't know how to touch this' branch of sampler replacement: unlike 626 there is no prior exception; the object simply fails an isinstance/duck-type check.

Source

Thrown at src/lightning/pytorch/utilities/data.py:317

                raise TypeError(
                    " Lightning can't inject a (distributed) sampler into your batch sampler, because it doesn't"
                    " subclass PyTorch's `BatchSampler`. To mitigate this, either follow the API of `BatchSampler` and"
                    " instantiate your custom batch sampler inside the `*_dataloader` hook of your module,"
                    " or set `Trainer(use_distributed_sampler=False)`. If you choose the latter, you will be"
                    " responsible for handling the distributed sampling within your batch sampler."
                ) from ex
        elif is_predicting:
            rank_zero_warn(
                f"You are using a custom batch sampler `{batch_sampler_cls.__qualname__}` for prediction."
                " Lightning would normally set `drop_last=False` to ensure all samples are returned, but for"
                " custom samplers it can't guarantee this. Make sure your sampler is configured correctly to return"
                " all indices.",
                category=PossibleUserWarning,
            )
        else:
            # The sampler is not a PyTorch `BatchSampler`, we don't know how to inject a custom sampler or
            # how to adjust the `drop_last` value
            raise TypeError(
                " Lightning can't inject a (distributed) sampler into your batch sampler, because it doesn't"
                " subclass PyTorch's `BatchSampler`. To mitigate this, either follow the API of `BatchSampler`"
                " or set `Trainer(use_distributed_sampler=False)`. If you choose the latter, you will be"
                " responsible for handling the distributed sampling within your batch sampler."
            )

        if is_predicting:
            batch_sampler = _IndexBatchSamplerWrapper(batch_sampler)

        # batch_sampler option is mutually exclusive with batch_size, shuffle, sampler, and drop_last
        return {
            "sampler": None,
            "shuffle": False,
            "batch_sampler": batch_sampler,
            "batch_size": 1,
            "drop_last": False,
        }

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Subclass torch.utils.data.sampler.BatchSampler so Lightning can re-instantiate it
  2. Or set Trainer(use_distributed_sampler=False) and handle sharding across ranks inside your sampler

Example fix

# before
loader = DataLoader(dataset, batch_sampler=MyCustomIterable())

# after
loader = DataLoader(dataset, batch_sampler=MyBatchSampler(sampler, batch_size=32, drop_last=False))
Defensive patterns

Strategy: type-guard

Validate before calling

from torch.utils.data import BatchSampler, DataLoader
assert isinstance(loader, DataLoader) and (loader.batch_sampler is None or isinstance(loader.batch_sampler, BatchSampler))

Type guard

from torch.utils.data import BatchSampler
from typing import Any

def batch_sampler_is_supported(obj: Any) -> bool:
    return obj is None or isinstance(obj, BatchSampler)

Prevention

When it happens

Trigger: Passing an arbitrary object as DataLoader(batch_sampler=...) that does not follow the BatchSampler API, during distributed sampler replacement triggered by DDP/FSDP-style strategies.

Common situations: Custom iterable batch samplers or generator-based samplers; single-device code moved to multi-GPU; mock/test objects used as batch samplers.

Related errors


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