fn

collate

Tensor or list or dict or tuple
collate(batch: list[object], collate_fn_map: dict[type, Callable[..., object]] | None = None)
source

Composable collate dispatcher with user-overridable type handlers.

Same dispatch behaviour as default_collate plus an optional collate_fn_map that overrides handling for specific element types. Use this when a custom record type needs special-case batching while everything else (tensors, dicts, namedtuples, ...) should follow the default rules.

Parameters

batchlist of object
Samples emitted by a dataset for a single mini-batch.
collate_fn_mapdict= None
Mapping from element type to a callable fn(batch, *, collate_fn_map=...). The first isinstance(batch[0], t) match wins; on miss the function falls back to default_collate's dispatch chain.

Returns

Tensor or list or dict or tuple

Collated batch. Structure mirrors default_collate's output.

Notes

When collate_fn_map is None this is exactly default_collate. The recursive forwarding of collate_fn_map lets nested containers reuse the same overrides.

Examples

>>> loader = DataLoader(ds, collate_fn=lambda b: collate(
...     b, collate_fn_map={MyRecord: my_record_collate}))