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

What it means

Raised as a TypeError (chained from an earlier TypeError) when Lightning attempts to rebuild a custom batch sampler by calling it with PyTorch BatchSampler-style arguments and the call fails in an unexpected way; because the class does not subclass torch.utils.data.sampler.BatchSampler, Lightning cannot safely inject a (distributed) sampler. The message lists two mitigations: follow the BatchSampler API inside a *_dataloader hook, or disable Lightning's sampler replacement.

Source

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

            # 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=(False if is_predicting else 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` 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(

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Make your batch sampler subclass torch.utils.data.sampler.BatchSampler (sampler, batch_size, drop_last attributes)
  2. Set Trainer(use_distributed_sampler=False) and implement distributed logic yourself
  3. Instantiate the custom batch sampler inside the *_dataloader hook so Lightning handles it

Example fix

# before
class MyBatchSampler:  # duck-typed, not a BatchSampler
    ...

# after
from torch.utils.data import BatchSampler
class MyBatchSampler(BatchSampler):
    ...
Defensive patterns

Strategy: type-guard

Validate before calling

from torch.utils.data import BatchSampler
assert isinstance(loader.batch_sampler, BatchSampler), 'batch_sampler must subclass BatchSampler for distributed runs'

Type guard

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

def is_batch_sampler(obj: Any) -> bool:
    return isinstance(obj, BatchSampler)

Prevention

When it happens

Trigger: DataLoader with a non-BatchSampler batch_sampler class whose reinstantiation with standard args raises an unexpected TypeError during distributed sampler replacement in _dataloader_init_kwargs_resolve_sampler.

Common situations: Custom batch sampler classes that accept incompatible positional args; versions of PyTorch where BatchSampler's signature changed; distributed training with elaborate custom samplers.

Related errors


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