Source code for lightkit.data.collation

from typing import Tuple
import torch


def collate_tuple(batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
    """
    Collate a tuple of batch items by returning the input tuple.

    This is the default used by
    :class:`~lightkit.data.DataLoader` when slices are cut from the underlying data source.
    """
    return batch


def collate_tensor(batch: Tuple[torch.Tensor, ...]) -> torch.Tensor:
    """
    Collates a tuple of batch items into the first tensor.

    Might be useful if only a single tensor is passed to
    :class:`~lightkit.data.DataLoader`.
    """
    return batch[0]