logging.basicConfig(level=logging.INFO)
logging.info("Starting model training...")

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

model.fit(train_dataset, epochs=5, callbacks=[tensorboard_callback])
model.evaluate(test_dataset)

logging.info("Model training completed.")

        __
        __