Lightning-AI/pytorch-lightning · warning

Trying to infer the `batch_size` from an ambiguous collectio

Error message

Trying to infer the `batch_size` from an ambiguous collection. The batch size we found is {batch_size}. To avoid any miscalculations, use `self.log(..., batch_size=batch_size)`.

What it means

PyTorch Lightning tries to automatically infer the batch size from the batch object when self.log(...) is called without an explicit batch_size argument. This warning fires when the batch is an ambiguous nested collection (dict/list/tuple of tensors) whose elements disagree about the batch size — e.g. a per-sample tensor alongside a [batch, ...] tensor — so the inferred number may be wrong. Metrics like accuracy or any normalization by batch size can then be miscalculated.

Source

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

def extract_batch_size(batch: BType) -> int:
    """Unpack a batch to find a ``torch.Tensor``.

    Returns:
        ``len(tensor)`` when found, or ``1`` when it hits an empty or non iterable.

    """
    error_msg = (
        "We could not infer the batch_size from the batch. Either simplify its structure"
        " or provide the batch_size as `self.log(..., batch_size=batch_size)`."
    )
    batch_size = None
    try:
        for bs in _extract_batch_size(batch):
            if batch_size is None:
                batch_size = bs
            elif batch_size != bs:
                warning_cache.warn(
                    "Trying to infer the `batch_size` from an ambiguous collection. The batch size we"
                    f" found is {batch_size}. To avoid any miscalculations, use `self.log(..., batch_size=batch_size)`."
                )
                break
    except RecursionError:
        raise RecursionError(error_msg)

    if batch_size is None:
        raise MisconfigurationException(error_msg)

    return batch_size


def has_len_all_ranks(
    dataloader: object,
    strategy: "pl.strategies.Strategy",
    allow_zero_length_dataloader_with_multiple_devices: bool = False,
) -> TypeGuard[Sized]:

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Pass the batch size explicitly at the log site: self.log('loss', loss, batch_size=batch['x'].size(0)).
2. Restructure the batch so all per-sample tensors share the same leading dimension and wrap non-per-sample metadata in a non-iterable object or keep it out of the batch dict.
3. Audit your batch composition by printing [_extract_batch_size(b) for b in one batch] in a sanity check to find which element disagrees.
4. If the inferred size is actually acceptable, pass batch_size explicitly anyway to make logging deterministic and silence the warning.

Example fix

# before
self.log("train_loss", loss, prog_bar=True)

# after
self.log("train_loss", loss, prog_bar=True, batch_size=x.size(0))

from lightning.pytorch.utilities.data import extract_batch_size

bs_values = set(extract_batch_size(batch))
assert len(bs_values) == 1, f"Ambiguous batch sizes in batch: {bs_values}"
self.log("loss", loss, batch_size=next(iter(bs_values)))

def has_uniform_batch_size(batch) -> bool:
    sizes = set(extract_batch_size(batch))
    return len(sizes) == 1

When it happens

Trigger: Calling self.log('name', value) (without batch_size=...) with a batch containing heterogeneous leading dimensions: e.g. {'x': tensor(B,3,224,224), 'idx': tensor(B)} mixed with a scalar/1-element structure, or nested dicts where one branch yields bs=N and another yields bs=1 or bs=N*seq_len. _extract_batch_size walks the collection, finds conflicting sizes, and warns (breaking after the first conflict).

Common situations: Custom batches with extra metadata tensors (indices, lengths, masks of shape [B] vs [1]); seq2seq batches where sequence tensors are reshaped to [B*T, ...]; token-level NLP batches mixing [B] input_ids with [B, T] flattened tensors; datasets returning tuples with a scalar label wrapped in a list.

Related errors

• you tried to log {v} which is currently not supported. Try a
• You are trying to `self.log()` but the loop's result collect
• You are trying to `self.log()` but it is not managed by the
• With `def training_step(self, dataloader_iter)`, `self.log(.
• ReduceLROnPlateau conditioned on metric {monitor_key} which

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