def main(
        gpu:Param("GPU to run on", str)=None,
        lr: Param("Learning rate", float)=1e-3,
        size: Param("Size (px: 128,192,224)", int)=128,
        debias_mom: Param("Debias statistics", bool)=False,
        debias_sqr: Param("Debias statistics", bool)=False,
        opt: Param("Optimizer: 'adam','genopt','rms','sgd'", str)='genopt',
        alpha: Param("Alpha", float)=0.99,
        mom: Param("Momentum", float)=0.9,
        eps: Param("epsilon", float)=1e-7,
        decay: Param("Decay AvgStatistic (momentum)", bool)=False,
        epochs: Param("Number of epochs", int)=5,
        bs: Param("Batch size", int)=128,
        ):
    """Distributed training of Imagenette.
    Fastest multi-gpu speed is if you run with: python -m fastai.launch"""

    # Pick one of these
    gpu = setup_distrib(gpu)
    if gpu is None: bs *= torch.cuda.device_count()

    moms = (mom,mom)
    stats = [
            AvgStatistic('momentum', mom,   scope=StatScope.Weight, decay=decay, debias=debias_mom),
            AvgSquare   ('alpha',    alpha, scope=StatScope.Weight, debias=debias_sqr),
            ConstStatistic('eps', eps), ConstStatistic('clip', 0.001),
            ]
    if   opt=='adam'  : opt_func = partial(optim.Adam, betas=(mom,alpha), eps=eps)
    elif opt=='rms'   : opt_func = partial(optim.RMSprop, alpha=alpha)
    elif opt=='genopt': opt_func = partial(GeneralOptimizer, on_step=on_step, stats=stats)
    else: raise Exception(f'unknown opt: {opt}')

    #opt_func = optim.SGD
    #learn = (cnn_learner(data, models.xresnet50, pretrained=False, concat_pool=False, lin_ftrs=[], split_on=bn_and_final,
    print(f'lr: {lr}; size: {size}; debias_mom: {debias_mom}; debias_sqr: {debias_sqr}; opt: {opt}; alpha: {alpha}; mom: {mom}; eps: {eps}; decay: {decay}')
    print('imagenette')
    get_learner(lr, size, False, bs, opt_func, gpu, epochs)
    gc.collect()

    print('imagewoof')
    get_learner(lr, size, True, bs, opt_func, gpu, epochs)