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`.setup_dataloaders(..., 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

Lightning attempted to re-instantiate a custom batch sampler to swap in a distributed sampler by calling it with the same signature as PyTorch's `BatchSampler.__init__(sampler, batch_size, drop_last, ...)`. The call raised a TypeError indicating the class doesn't follow that API (it isn't a `BatchSampler` subclass), so Lightning cannot inject the sampler and re-raises with guidance.

Source

Thrown at src/lightning/fabric/utilities/data.py:220

            # This is a sampler for which we could not capture the init args, but it kinda looks like a batch sampler
            # even if it does not inherit from PyTorch's interface.
            try:
                batch_sampler = batch_sampler_cls(
                    sampler,
                    batch_size=batch_sampler.batch_size,
                    drop_last=batch_sampler.drop_last,
                )
            except TypeError as ex:
                import re

                match = re.match(r".*__init__\(\) (got multiple values)|(missing \d required)", str(ex))
                if not match:
                    # an unexpected `TypeError`, continue failure
                    raise

                # There could either be too few or too many arguments. Customizing the message based on this doesn't
                # make much sense since our MisconfigurationException is going to be raised from the original one.
                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`.setup_dataloaders(..., use_distributed_sampler=False)`. If you choose the latter, you"
                    " will be responsible for handling the distributed sampling within your batch sampler."
                ) from ex
        else:
            # The sampler is not a PyTorch `BatchSampler`, we don't know how to inject a custom sampler
            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`.setup_dataloaders(..., use_distributed_sampler=False)`. If you choose the latter, you"
                " will be responsible for handling the distributed sampling within your batch sampler."
            )

        return {
            "sampler": None,
            "shuffle": False,
            "batch_sampler": batch_sampler,

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Set `use_distributed_sampler=False` in `setup_dataloaders(...)` and implement distributed sharding inside your batch sampler yourself.
  2. Refactor your custom class to subclass `torch.utils.data.BatchSampler` and follow its `__init__(sampler, batch_size, drop_last)` API.
  3. Use a regular `sampler=` DataLoader instead of `batch_sampler=` so Lightning can wrap it with a DistributedSampler.

Example fix

# before
loader = DataLoader(dataset, batch_sampler=MyCustomBatchSampler(...))
fabric.setup_dataloaders(loader)

# after
loader = DataLoader(dataset, batch_sampler=MyCustomBatchSampler(...))
fabric.setup_dataloaders(loader, use_distributed_sampler=False)
Defensive patterns

Strategy: fallback

Validate before calling

from torch.utils.data import BatchSampler

def is_standard_batch_sampler(bs) -> bool:
    return isinstance(bs, BatchSampler)

Prevention

When it happens

Trigger: A DataLoader with a `batch_sampler` that is not a subclass of `torch.utils.data.BatchSampler` is passed to Fabric's `setup_dataloaders` while distributed sampling is enabled; Lightning tries `type(batch_sampler)(sampler=..., batch_size=..., drop_last=...)` and the constructor rejects those arguments.

Common situations: Custom iterable batch sampler classes (e.g. grouping buckets, weighted batch samplers) used with multi-GPU DDP training in Lightning Fabric.

Related errors


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