def get_input_shape(dataset: str):
    """Return a list of integer indicating the input shape as (num_channel, height, weight)"""
    if dataset == "imagenet":
        return (3, 224, 224)
    elif dataset == 'cifar10':
        return (3, 32, 32)
    elif dataset == 'mnist':
        return (1, 28, 28)