import sys

from collect_generalizability_metrics import get_generalizability_gap

from training.vanilla_trainer import VanillaTrainer

if __name__ == '__main__':

    config = {"model": "DeepLab",

              "device": "cuda",

              "lr": 0.00001,

              "batch_size": 8,

              "epochs": 250}

    trainer = VanillaTrainer(sys.argv[1], config)

    trainer.train()

    # for i in ["TriUnet", "DeepLab", "FPN", "Unet"]:

    #     config["model"] = i

    #     trainer = VanillaTrainer(sys.argv[1], config)

    #     trainer.train()

    # for i in range(13, 100):

    #     trainer = VanillaTrainer("DeepLab", i, config)

    #     trainer.train()

    # i = 4

    # while True:

    #     config = {"epochs": 200, "id": i, "lr": 0.00001, "pretrain": False}

    #     training.vanilla_trainer.train_vanilla_predictor(config)

    #     i += 1