This is what I just used:
` import os
from argparse import ArgumentParser

from tensorflow.keras import utils as np_utils
from tensorflow.keras.datasets import mnist
from tensorflow.keras.layers import Activation, Dense, Softmax
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ModelCheckpoint

from clearml import Task

parser = ArgumentParser()
parser.add_argument('--output-uri', type=str, required=False)
args = parser.parse_args()

the data, shuffled and split between train and test sets

nb_classes = 10
(X_train, y_train), (X_test, y_test) = mnist.load_data()

X_train = X_train.reshape(60000, 784).astype('float32')/255.
X_test = X_test.reshape(10000, 784).astype('float32')/255.
print(X_train.shape[0], 'train samples')
print(X_test.shape[0], 'test samples')

convert class vectors to binary class matrices

Y_train = np_utils.to_categorical(y_train, nb_classes)
Y_test = np_utils.to_categorical(y_test, nb_classes)

model = Sequential()
model.add(Dense(10, input_shape=(784,)))
model.add(Activation('relu'))
model.add(Softmax())
model.summary()
output_folder = os.path.join(os.getcwd(), 'tmp')
model.compile(loss='categorical_crossentropy',
optimizer=Adam(),
metrics=['accuracy'])

model_checkpoint = ModelCheckpoint('best_model', save_best_only=True)

Connecting ClearML

task = Task.init(project_name='examples',
task_name='Upload problem',
output_uri=args.output_uri)

history = model.fit(X_train, Y_train,
batch_size=128,
epochs=5,
callbacks=[model_checkpoint],
verbose=1,
validation_data=(X_test, Y_test))

os.makedirs(output_folder, exist_ok=True)
model.save(os.path.join(output_folder, 'model.h5'))

print('Number of output models: {}'.format(len(task.models["output"]))) `