def guess_data_type(shape, risky=False):
    """Infer the type of data based on the shape of the tensors

    Args:
        risky(bool): some guesses are more likely to be wrong.
    """
    if len(shape) in (1, 2):
        return 'label'
    if risky and len(shape) == 3:
        return 'image'
    if len(shape) == 4:
        if shape[-1] in (1, 3, 4):
            return 'image'
        else:
            return 'segmentation_mask'
    return None
