def collate_fn(batch):
        """
        Collect batch for YOLO X model.

        Args:
            batch (List[Dict[str, torch.Tensor]]):
                List with records from YOLOXDataset.

        Returns:
            images batch with shape [B, C, H, W]
            boxes with shape [B, MAX_OBJECTS, 4]
            classes with shape [B, MAX_OBJECTS]
        """
        images, boxes, classes = [], [], []
        for item in batch:
            images.append(item["image"])
            boxes.append(item["boxes"])
            classes.append(item["classes"])

        images = torch.stack(images)
        boxes = torch.stack(boxes)
        classes = torch.stack(classes)

        return {"image": images, "bboxes": boxes, "labels": classes}

        Collect batch for YOLO X model.

        Args:
            batch (List[Dict[str, torch.Tensor]]):
                List with records from YOLOXDataset.

        Returns:
            images batch with shape [B, C, H, W]
            boxes with shape [B, MAX_OBJECTS, 4]
            classes with shape [B, MAX_OBJECTS]
        
examples/detection/dataset.py