Lightning-AI/pytorch-lightning · error · TypeError

Trying to inject a modified sampler into the batch sampler;

Error message

Trying to inject a modified sampler into the batch sampler; however, it seems the class `{batch_sampler_cls.__qualname__}` does not have an argument called `sampler.` To mitigate this, expose an argument `sampler` in the `__init__` method of your custom class.

What it means

When Lightning replaces the sampler inside a custom batch sampler during distributed setup, it re-instantiates the batch sampler class with a new (distributed) sampler via the `sampler` argument. If the custom batch sampler class (which implemented `__pl_saved_arg_names__`) doesn't declare a `sampler` parameter in its `__init__`, the replacement cannot be performed and this TypeError is raised.

Source

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

) -> dict[str, Any]:
    """This function is used to handle the sampler, batch_sampler arguments associated within a DataLoader for its re-
    instantiation."""
    batch_sampler = getattr(dataloader, "batch_sampler")

    if batch_sampler is not None and type(batch_sampler) is not BatchSampler:
        batch_sampler_cls = type(batch_sampler)
        if hasattr(batch_sampler, "__pl_saved_args"):
            # This is a PyTorch `BatchSampler` subclass for which we captured the init args
            args = batch_sampler.__pl_saved_args
            kwargs = batch_sampler.__pl_saved_kwargs
            default_kwargs = batch_sampler.__pl_saved_default_kwargs
            arg_names = batch_sampler.__pl_saved_arg_names

            success, args, kwargs = _replace_value_in_saved_args(
                "sampler", sampler, args, kwargs, default_kwargs, arg_names
            )
            if not success:
                raise TypeError(
                    "Trying to inject a modified sampler into the batch sampler; however, it seems the class "
                    f"`{batch_sampler_cls.__qualname__}` does not have an argument called `sampler.` To mitigate "
                    "this, expose an argument `sampler` in the `__init__` method of your custom class."
                )

            batch_sampler = _reinstantiate_wrapped_cls(batch_sampler, *args, **kwargs)
        elif hasattr(batch_sampler, "batch_size") and hasattr(batch_sampler, "drop_last"):
            # 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

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Add a `sampler` argument to your custom batch sampler's `__init__` and use it instead of constructing an internal sampler.
  2. Set `use_distributed_sampler=False` in `setup_dataloaders(...)` and handle distributed sampling yourself.
  3. Use PyTorch's `BatchSampler` (or subclass it) so Lightning knows its API.

Example fix

// before
class MyBatchSampler:
    def __init__(self, sampler, batch_size, drop_last):
        self.sampler = MyOwnSampler(...)

// after
class MyBatchSampler:
    def __init__(self, sampler, batch_size, drop_last):
        self.sampler = sampler  # accept injected (distributed) sampler
Defensive patterns

Strategy: validation

Validate before calling

import inspect

def batch_sampler_injectable(batch_sampler) -> bool:
    params = inspect.signature(type(batch_sampler).__init__).parameters
    return "sampler" in params or any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())

Prevention

When it happens

Trigger: Passing a DataLoader with a custom `batch_sampler` whose class records saved arg names (via pickling hooks) but whose `__init__` has no `sampler` parameter, then running Fabric's `setup_dataloaders` (with distributed sampling enabled) so Lightning tries to inject a distributed sampler into the batch sampler.

Common situations: Custom BatchSampler implementations that take a dataset + indices but construct their own internal sampler instead of receiving one via `__init__`; multi-GPU/DDP runs where sampler injection is mandatory.

Related errors

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