def get_task_dimensions(task):
    """
    Get the number of channels, channel size (height * width) and number of classes of a dataset.
    :param task: str, '*MNIST' or 'CIFAR*' for any string *
    :return: int, int, int; in_channels, input_size, n_classes
    """
    if 'MNIST' in task:
        in_channels = 1
        input_size = 28 * 28
        n_classes = 10
    elif 'CIFAR' in task:
        in_channels = 3
        input_size = 32 * 32
        n_classes = int(task[5:])
    else:
        raise ValueError('Task must be either *MNIST or CIFAR*, but %s was given' % task)
    return in_channels, input_size, n_classes