"""Stage-2: train LED initializer on SDD (baseline or +graph).""" import argparse from trainer.train_sdd_led import Trainer def main(): p = argparse.ArgumentParser() p.add_argument('--cuda', default=True) p.add_argument('--cfg', default='sdd/sdd') p.add_argument('--gpu', type=int, default=0) p.add_argument('--info', type=str, default='baseline') p.add_argument('--learning_rate', type=float, default=0.002) p.add_argument('--train', type=int, default=1) p.add_argument('--grad_accum', type=int, default=16) p.add_argument('--use_graph', action='store_true') p.add_argument('--use_v6_graph', action='store_true') p.add_argument('--residual_on', type=str, default='y0') config = p.parse_args() Trainer(config).fit() if __name__ == '__main__': main()