import torch


def collate_test(batch):
    """
    Collate function for preparing the testing batch for DCL testing during training.

    :param batch: The list containing outputs of the __get_item__() function.
    Length of the list is equal to the required batch size.
    :return: The batch containing images and actual labels for testing
    """
    imgs = []  # List to store the images
    target = []  # List to store the targets
    # Iterate over each output of the __get_item__() function
    for sample in batch:
        imgs.append(sample[0])  # Append the original image
        target.append(sample[1])  # Append the origin class label
    # Stack the images and return the test batch
    return torch.stack(imgs, 0), target