05b2b1c02a8ce5f471877243ebebe5df876ccc37,models/unet/train_unet.py,,main,#,194

Before Change

def main(args):
    if args.mode == "train":
        load_version = 0 if args.resume else None
        logger = TestTubeLogger(save_dir=args.exp_dir, name=args.exp, version=load_version)
        trainer = create_trainer(args, logger)
        model = UnetMRIModel(args)
        trainer.fit(model)
    else:  // args.mode == "test"
        assert args.checkpoint is not None
        model = UnetMRIModel.load_from_checkpoint(str(args.checkpoint))
        model.hparams.sample_rate = 1.
        trainer = create_trainer(args, logger=False)
        trainer.test(model)


if __name__ == "__main__":
    parser = Args()
    parser.add_argument("--mode", choices=["train", "test"], default="train")
    parser.add_argument("--num-epochs", type=int, default=50, help="Number of training epochs")

After Change

    if args is not None:
        parser.set_defaults(**args)

    args, _ = parser.parse_known_args()
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    run(args)

Instances