def train_models():
    wandb.init()
    config = wandb.config
    strategy = tf.distribute.MirroredStrategy()
    print(f"Number of GPUS : {strategy.num_replicas_in_sync}")

    with strategy.scope():
        VGG_model = VGG((config.IMG_SIZE,config.IMG_SIZE,3),len(class_names),
                       num_blocks = config.num_blocks,filter_mult=config.filter_mults,
                        dense_neurons=config.units,dropout=config.dropouts)
        VGG_model(tf.keras.layers.Input(shape=(config.IMG_SIZE,config.IMG_SIZE,3)))
        early_stopping_cb = tf.keras.callbacks.EarlyStopping(patience=3)
        VGG_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=config.learning_rate),
                      loss=tf.keras.losses.CategoricalCrossentropy(),
                      metrics=["accuracy"])

        history = VGG_model.fit(
            train_dataset,
            epochs=config.epochs,
            validation_data=valid_dataset,
            callbacks=[early_stopping_cb, 
                       WandbMetricsLogger(log_freq=10)]
        )
        val_loss = history.history["val_loss"][-1]
        train_loss = history.history["loss"][-1]
        val_accuracy = history.history["val_accuracy"][-1]
        accuracy = history.history["accuracy"][-1]
        wandb.log({"val_loss": val_loss, 
                   "train_loss": train_loss,
                  "val_accuracy":val_accuracy,
                  "accuracy":accuracy})

        __
        __