overfit_batches replaces custom sampler with Sequential sampler

Nobody has claimed this yet.

### Bug description

More complicated datasets for real-world applications can demand more sophisticated data sampling, that are not just indices but perhaps a dictionary with additional info to query the dataset and return samples. When using the overfit_batches setting, which should still be useful in these settings to test an implementation, the dataloaders fail, because in the overfit_batches setting any custom sampler passed to the dataset loader is replaced with a standard SequentialLoader from pytorch.

An example of a domain library that extensively makes use of custom query logic for geospatial data and is heavily integrated into lightning is [torchgeo](https://github.com/torchgeo/torchgeo).

### What version are you seeing the problem on?

v2.5

### Reproduced in studio

_No response_

### How to reproduce the bug

```python
"""Minimal reproducible example: Custom dict sampler with Lightning overfit_batches."""

import torch
from torch.utils.data import Dataset, DataLoader, Sampler
from lightning import LightningModule, LightningDataModule, Trainer
from collections.abc import Iterator
import numpy as np

# ============================================================================
# 1. SIMPLE DATASET (accepts dict queries)
# ============================================================================

class SimpleQueryDataset(Dataset):
"""Simple dataset that accepts dictionary queries."""

def __init__(self, n_items: int = 100):
"""Initialize with some dummy items."""
self.n_items = n_items

def __len__(self) -> int:
return self.n_items

def __getitem__(self, query: dict) -> dict:
"""Accept a dictionary query.

Args:
query: Dict with keys:
- 'index': int identifier
- 'param': some parameter value

Returns:
Dict with random dummy data
"""
# Validate input type
if not isinstance(query, dict):
raise TypeError(
f"Expected dict query, got {type(query)}. "
f"Content: {query}"
)

index = query["index"]
param = query["param"]

# Generate deterministic data based on query
n_points = 10 + (index % 20)

return {
"data": torch.randn(n_points, 3),
"target": torch.randn(n_points, 1),
"query_index": index,
"query_param": param,
}

# ============================================================================
# 2. SIMPLE QUERY SAMPLER (yields dict queries)
# ============================================================================

class SimpleQuerySampler(Sampler):
"""Simple sampler that yields dictionary queries."""

def __init__(self, n_queries: int = 50, shuffle: bool = True, seed: int = 42):
"""Generate some simple queries."""
self.n_queries = n_queries
self.shuffle = shuffle
self.seed = seed
self.rng = np.random.RandomState(seed)

# Pre-generate queries
self.queries = []
for i in range(n_queries):
self.queries.append({
"index": i,
"param": self.rng.uniform(0, 1),
})

def __iter__(self) -> Iterator[dict]:
"""Yield query dictionaries."""
indices = list(range(self.n_queries))

if self.shuffle:
self.rng.shuffle(indices)

for idx in indices:
yield self.queries[idx]

def __len__(self) -> int:
return self.n_queries

# ============================================================================
# 3. LIGHTNING DATAMODULE
# ============================================================================

class SimpleDataModule(LightningDataModule):
"""Lightning DataModule using dict-based query sampler."""

def __init__(self, batch_size: int = 4, num_workers: int = 0):
super().__init__()
self.batch_size = batch_size
self.num_workers = num_workers

# Create dataset
self.dataset = SimpleQueryDataset(n_items=100)

def train_dataloader(self) -> DataLoader:
sampler = SimpleQuerySampler(n_queries=50, shuffle=True, seed=42)

return DataLoader(
self.dataset,
batch_size=self.batch_size,
sampler=sampler,
num_workers=self.num_workers,
collate_fn=self._collate_fn,
)

def val_dataloader(self) -> DataLoader:
sampler = SimpleQuerySampler(n_queries=20, shuffle=False, seed=43)

return DataLoader(
self.dataset,
batch_size=self.batch_size,
sampler=sampler,
num_workers=self.num_workers,
collate_fn=self._collate_fn,
)

def _collate_fn(self, batch: list[dict]) -> dict:
"""Simple collate function."""
# Stack data across batch
data = torch.cat([sample["data"] for sample in batch], dim=0)
target = torch.cat([sample["target"] for sample in batch], dim=0)

return {
"data": data,
"target": target,
"batch_sizes": [sample["data"].shape[0] for sample in batch],
"query_indices": [sample["query_index"] for sample in batch],
}

# ============================================================================
# 4. SIMPLE MODEL
# ============================================================================

class SimpleModel(LightningModule):
"""Simple model for testing."""

def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(3, 1)

def forward(self, x):
return self.linear(x)

def training_step(self, batch, batch_idx):
data = batch["data"]
target = batch["target"]
pred = self(data)
loss = torch.nn.functional.mse_loss(pred, target)

self.log("train_loss", loss)
print(f" Train batch {batch_idx}: query_indices={batch['query_indices']}")
return loss

def validation_step(self, batch, batch_idx):
data = batch["data"]
target = batch["target"]
pred = self(data)
loss = torch.nn.functional.mse_loss(pred, target)

self.log("val_loss", loss)
print(f" Val batch {batch_idx}: query_indices={batch['query_indices']}")
return loss

def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=1e-3)

# ============================================================================
# TEST FUNCTIONS
# ============================================================================

def test_lightning_overfit():
"""Test: Lightning WITH overfit_batches (THE FAILING CASE)."""
datamodule = SimpleDataModule(batch_size=2, num_workers=0)
model = SimpleModel()

trainer = Trainer(
max_epochs=1,
accelerator="cpu",
enable_checkpointing=False,
logger=False,
enable_model_summary=False,
overfit_batches=2, # This causes the issue
)

trainer.fit(model, datamodule)

if __name__ == "__main__":
test_lightning_overfit()
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```

### More info

_No response_

cc @tchaton

Start by running the minimal reproducible example with Trainer configured with overfit_batches and inspect how the train and validation DataLoaders are modified. Trace the overfit_batches handling and the replacement of the custom sampler with SequentialSampler. Done means the provided dict-yielding sampler remains usable and the example completes for both loaders under overfit_batches.

Written by the indexing model from the issue text.