def _get_new_batch_shape(param: ParamItem, batch_shape: torch.Size) -> torch.Size:
    """Get the new batch shape if the augmentation changes the image size.

    Note:
       Augmentations that change the image size must provide the parameter `output_size`.

    """
    if param.data is None:
        return batch_shape
    if isinstance(param.data, list):
        for p in param.data:
            batch_shape = _get_new_batch_shape(p, batch_shape)
    elif "output_size" in param.data:
        if not (param.data["batch_prob"] > 0.5)[0]:
            # Augmentations that change the image size must be applied equally to all elements in batch.
            # If the augmentation is not applied, return the same batch shape.
            return batch_shape
        new_batch_shape = list(batch_shape)
        new_batch_shape[-2:] = param.data["output_size"][0]
        batch_shape = torch.Size(new_batch_shape)
    return batch_shape
Get the new batch shape if the augmentation changes the image size.

    Note:
       Augmentations that change the image size must provide the parameter `output_size`.

    
kornia/augmentation/container/image.py