{"record":{"id":"add9b557c4af69aa","repo":"Lightning-AI/pytorch-lightning","slug":"trying-to-inject-a-modified-sampler-into-the-batch","errorCode":null,"errorMessage":"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.","messagePattern":"Trying to inject a modified sampler into the batch sampler; however, it seems the class `(.+?)` does not have an argument called `sampler\\.` To mitigate this, expose an argument `sampler` in the `__init__` method of your custom class\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"src/lightning/fabric/utilities/data.py","lineNumber":194,"sourceCode":") -> dict[str, Any]:\n    \"\"\"This function is used to handle the sampler, batch_sampler arguments associated within a DataLoader for its re-\n    instantiation.\"\"\"\n    batch_sampler = getattr(dataloader, \"batch_sampler\")\n\n    if batch_sampler is not None and type(batch_sampler) is not BatchSampler:\n        batch_sampler_cls = type(batch_sampler)\n        if hasattr(batch_sampler, \"__pl_saved_args\"):\n            # This is a PyTorch `BatchSampler` subclass for which we captured the init args\n            args = batch_sampler.__pl_saved_args\n            kwargs = batch_sampler.__pl_saved_kwargs\n            default_kwargs = batch_sampler.__pl_saved_default_kwargs\n            arg_names = batch_sampler.__pl_saved_arg_names\n\n            success, args, kwargs = _replace_value_in_saved_args(\n                \"sampler\", sampler, args, kwargs, default_kwargs, arg_names\n            )\n            if not success:\n                raise TypeError(\n                    \"Trying to inject a modified sampler into the batch sampler; however, it seems the class \"\n                    f\"`{batch_sampler_cls.__qualname__}` does not have an argument called `sampler.` To mitigate \"\n                    \"this, expose an argument `sampler` in the `__init__` method of your custom class.\"\n                )\n\n            batch_sampler = _reinstantiate_wrapped_cls(batch_sampler, *args, **kwargs)\n        elif hasattr(batch_sampler, \"batch_size\") and hasattr(batch_sampler, \"drop_last\"):\n            # This is a sampler for which we could not capture the init args, but it kinda looks like a batch sampler\n            # even if it does not inherit from PyTorch's interface.\n            try:\n                batch_sampler = batch_sampler_cls(\n                    sampler,\n                    batch_size=batch_sampler.batch_size,\n                    drop_last=batch_sampler.drop_last,\n                )\n            except TypeError as ex:\n                import re\n","sourceCodeStart":176,"sourceCodeEnd":212,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/utilities/data.py#L176-L212","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Add a `sampler` argument to your custom batch sampler's `__init__` and use it instead of constructing an internal sampler.","Set `use_distributed_sampler=False` in `setup_dataloaders(...)` and handle distributed sampling yourself.","Use PyTorch's `BatchSampler` (or subclass it) so Lightning knows its API."],"exampleFix":"// before\nclass MyBatchSampler:\n    def __init__(self, sampler, batch_size, drop_last):\n        self.sampler = MyOwnSampler(...)\n\n// after\nclass MyBatchSampler:\n    def __init__(self, sampler, batch_size, drop_last):\n        self.sampler = sampler  # accept injected (distributed) sampler","handlingStrategy":"validation","validationCode":"import inspect\n\ndef batch_sampler_injectable(batch_sampler) -> bool:\n    params = inspect.signature(type(batch_sampler).__init__).parameters\n    return \"sampler\" in params or any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Give custom batch samplers a `sampler` init argument.","If you can't, plan to call setup_dataloaders(use_distributed_sampler=False)."],"tags":["pytorch-lightning","batch-sampler","distributed","ddp"],"backgroundTag":"sampler-injection-unsupported","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}