import torch


def collate_fn(data):
    """Creates mini-batch tensors from the list of images.
    
    We should build custom collate_fn rather than using default collate_fn, 
    because merging caption (including padding) is not supported in default.
    Args:
        data: list
            - image: torch tensor of shape (3, 256, 256).
            
    Returns:
        images: torch tensor of shape (batch_size, 3, 256, 256).
        
    """

    comp_imgs, bg_imgs = zip(*data)
    comp_imgs = list(comp_imgs)
    bg_imgs = list(bg_imgs)
    # hole_images = list(hole_images)
    # masks = list(masks)

    return torch.stack(comp_imgs, dim=0), torch.stack(bg_imgs,0)