def collate_fn(batch):
    """
    Collates a batch of data for object detection training.

    Args:
        batch (list): A list of dictionaries, where each dictionary contains
                      'pixel_values', 'labels', and optionally 'pixel_mask'.

    Returns:
        dict: A dictionary with the following keys:
            - 'pixel_values' (torch.Tensor): A tensor containing stacked pixel values from the batch.
            - 'labels' (list): A list of labels from the batch.
            - 'pixel_mask' (torch.Tensor, optional): A tensor containing stacked pixel masks from the batch,
                                                     if 'pixel_mask' is present in the input batch.
    """
    data = {}
    data["pixel_values"] = torch.stack([x["pixel_values"] for x in batch])
    data["labels"] = [x["labels"] for x in batch]
    if "pixel_mask" in batch[0]:
        data["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch])
    return data

    Collates a batch of data for object detection training.

    Args:
        batch (list): A list of dictionaries, where each dictionary contains
                      'pixel_values', 'labels', and optionally 'pixel_mask'.

    Returns:
        dict: A dictionary with the following keys:
            - 'pixel_values' (torch.Tensor): A tensor containing stacked pixel values from the batch.
            - 'labels' (list): A list of labels from the batch.
            - 'pixel_mask' (torch.Tensor, optional): A tensor containing stacked pixel masks from the batch,
                                                     if 'pixel_mask' is present in the input batch.
    
src/autotrain/trainers/object_detection/utils.py