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

Raised as a TypeError when Lightning tries to replace the sampler inside a custom batch sampler (the batch_sampler argument of a DataLoader) but the batch sampler class's __init__ has no 'sampler' parameter to substitute. Lightning reinstantiates the batch sampler class with a modified sampler; without a 'sampler' init arg there is no way to inject the distributed one.

Source

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

            arg_names = batch_sampler.__pl_saved_arg_names

            if is_predicting:
                success, args, kwargs = _replace_value_in_saved_args(
                    "drop_last", False, args, kwargs, default_kwargs, arg_names
                )
                if not success:
                    rank_zero_warn(
                        f"Trying to inject `drop_last=False` into batch sampler since you are predicting, however "
                        f"it seems the class `{batch_sampler_cls.__qualname__}` does not support it. "
                        "Your predictions might be incomplete. To mitigate this, expose `drop_last` in "
                        "the `__init__` method of your custom class."
                    )

            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=(False if is_predicting else batch_sampler.drop_last),
                )
            except TypeError as ex:
                import re

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Expose a sampler argument in your batch sampler's __init__ and use it to build batches
  2. Or set Trainer(use_distributed_sampler=False) and handle distributed sampling inside your batch sampler yourself

Example fix

# before
class MyBatchSampler:
    def __init__(self, dataset, batch_size): ...

# after
class MyBatchSampler:
    def __init__(self, sampler, batch_size): ...  # 'sampler' arg injectable
Defensive patterns

Strategy: validation

Validate before calling

import inspect
bs = loader.batch_sampler
assert 'sampler' in inspect.signature(type(bs).__init__).parameters, 'batch sampler needs a sampler __init__ arg'

Type guard

import inspect
from typing import Any

def batch_sampler_has_sampler_arg(bs: Any) -> bool:
    return 'sampler' in inspect.signature(type(bs).__init__).parameters

Prevention

When it happens

Trigger: Passing a DataLoader with a custom batch_sampler whose class's __init__ lacks a sampler argument, in a distributed run where Lightning replaces samplers (e.g. FSDP/DDP with use_distributed_sampler default).

Common situations: Custom batch samplers with signatures like __init__(self, indices, batch_size); porting single-device code to multi-device training; using weighted or curriculum batch samplers.

Related errors

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