import torch


def boolean_mask(img, color):
    """
    Returns a Boolean mask on a image, based on the presence of a color.

    Arguments:
        img {torch.Tensor} -- image tensor [shape = (..., 3)]
        color {torch.Tensor} -- RGB color tensor [shape = (3, )]

    Returns:
        torch.BoolTensor -- boolean mask of image [shape = (..., )]
    """

    dim = len(img.shape) - 1

    return torch.all(img == color.view(*([1] * dim), 3), dim=dim)