d750d66fb998a44316000908f213ae73c53de97f,examples/shrec17/train.py,,,#,136

Before Change

        torch.save(model.state_dict(), "state.pkl")


if __name__ == "__main__":
    main()

After Change

        torch.save(model.state_dict(), os.path.join(log_dir, "state.pkl"))


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument("--log_dir", type=str, required=True)
    parser.add_argument("--augmentation", type=int, default=1,
                        help="Generate multiple image with random rotations and translations (recommanded = 3)")
    parser.add_argument("--dataset", choices={"test", "val", "train"}, default="train")
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--num_workers", type=int, default=1)

    args = parser.parse_args()

    main(**args.__dict__)

Instances