config = {
    "BATCH_SIZE" :32,
    "IMG_SIZE":150,
    "VAL_SPLIT":0.1,
    "EPOCHS":100,
    "PATIENCE":3,
    "lr":0.0001
}

def train_model(build_model,model_type):
    wandb.init(
        project="resnet-comparison",
        group="ResNet Variants",
        name=f"ResNet-{model_type}",  # Example: ResNet-18, ResNet-50, etc.
        config=config
    )
    strategy = tf.distribute.MirroredStrategy()
    early_stopping_cb = tf.keras.callbacks.EarlyStopping(patience=config["PATIENCE"])
    wandb_callback = wandb.keras.WandbCallback(save_model=False)
    with strategy.scope():
        model = build_model(num_classes=len(class_names))
        model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=config["lr"]),
                        loss=tf.keras.losses.CategoricalCrossentropy(),
                        metrics=["accuracy"])
        history = model.fit(
            train_dataset,
            epochs=config["EPOCHS"],
            validation_data=valid_dataset,
            callbacks=[early_stopping_cb,wandb_callback]
        )

    test_loss, test_acc = model.evaluate(test_dataset)
    wandb.log({
    "test loss": test_loss,
    "test Accuracy": test_acc
})

    wandb.finish()

        __
        __