import torch


def collate_fn(videos):
    """
    Collate function for the PyTorch data loader.
    Merges all batch videos in a tensor with shape (length, batch, channels, width, height) and converts their pixel
    values to [0, 1].
    Parameters
    ----------
    videos : list
        List of uint8 NumPy arrays representing videos with shape (length, batch, width, height, channels).
    Returns
    -------
    torch.*.Tensor
        Batch of videos with shape (length, batch, channels, width, height) and float values lying in [0, 1].
    """
    
    seq_len = len(videos[0])
    batch_size = len(videos)
    nc = 1 if videos[0].ndim == 3 else 3
    w = videos[0].shape[3]
    h = videos[0].shape[2]
    tensor = torch.zeros((seq_len, batch_size, nc, h, w))
    for i, video in enumerate(videos):
        if nc == 1:
            tensor[:, i, 0] += video
        if nc == 3:
            tensor[:, i] += video
    return tensor