def get_input_size(args):
    """
    Return the input size for the given arguments.    
    """
    if args.dataset in ('cifar10', 'cifar100'):
        return (1, 3, 32, 32)
    else:
        return (1, 3, 224, 224)