import torch


def center_crop(data, shape, include_batch: bool = False) -> torch.Tensor:
    """Apply a center crop to a batch of images.

    Args:
        data (torch.Tensor): The input tensor to be center cropped.
        shape (Tuple[int]): The output shape. If shape[dim] = -1 or None,
            then no crop will be applied in that dimension.
        include_batch (bool): If ``True``, the center crop will also be
            applied to the batch dimension (``dim=0``).

    Returns:
        torch.Tensor: The center-cropped tensor.

    Note:
        The 0-th dimension is assumed to be the batch dimension.
        To apply the crop to this dimension as well, set ``include_batch=True``.
    """
    for i in range(len(shape)):
        data_dim = i if include_batch else i + 1
        if (shape[i] == data.shape[data_dim]) or (shape[i] == -1):
            continue
        assert 0 < shape[i] <= data.shape[data_dim]
        idx_start = (data.shape[data_dim] - shape[i]) // 2
        data = data.narrow(data_dim, idx_start, shape[i])

    return data