LIST_MODEL = {
    "resnet18": build_resnet18,
    "resnet34": build_resnet34,
    "resnet50": build_resnet50,
    "resnet101": build_resnet101,
    "resnet152": build_resnet152,
}

# Train all models
for model_type, model in LIST_MODEL.items():
    train_model(model, model_type)

        __
        __