def parse_args(parser):
    parser.add_argument('--aug', action='store_true', default=False, help='data augmentation')
    parser.add_argument('--train_bs', default=2, type=int, help='batch size for training')
    parser.add_argument('--test_bs', default=2, type=int, help='batch size for test')
    parser.add_argument('--epochs', default=200, type=int, help='training epochs')