def test_collate_heterogeneous_shapes() -> None:
    """Test collating items with different shapes."""
    items = [
        DummyDatasetItem(
            image=Image(torch.rand((3, 256, 256))),
            mask=Mask(torch.ones((256, 256))),
        ),
        DummyDatasetItem(
            image=Image(torch.rand((3, 224, 224))),
            mask=Mask(torch.ones((224, 224))),
        ),
    ]
    batch = DummyBatch.collate(items)
    # the collated batch should have the shape of the largest item
    assert batch.image.shape == (2, 3, 256, 256)
Test collating items with different shapes.
tests/unit/data/dataclasses/test_collate.py