import torch


def collate_fn(batch):
    """
    Args:
        batch - [(xyz, feat, label, ...), ...] - each tuple from a sampler
    Returns: [
        xyz     : [BxN, 3]
        feat    : [BxN, d]
        label   : [BxN]
        ...
        offset  : int
    ]
    """
    batch_list = []
    for sample in batch:
        if isinstance(sample, list):
            batch_list += sample
        else:
            batch_list.append(sample)
    batch_list = list(zip(*batch_list))  # [[xyz, ...], [feat, ...], ...]

    offset, count = [], 0
    for item in batch_list[0]:
        count += item.shape[0]
        offset.append(count)
    offset = torch.IntTensor(offset)
    batch_list = [torch.cat(v) for v in batch_list]
    return [*batch_list, offset]