def __make_collates(input_transform: InputTransform, on_device: bool, collate: Callable) -> Tuple[Callable, Callable]:
    """Returns the appropriate collate functions based on whether the transforms happen in a DataLoader worker or on the
    device (main process)."""
    if on_device:
        return input_transform._identity, collate
    return collate, input_transform._identity
Returns the appropriate collate functions based on whether the transforms happen in a DataLoader worker or on the
    device (main process).
src/flash/core/data/io/input_transform.py