mlflow save model

lab5/train/Jenkinsfile

@@ -46,6 +46,7 @@ pipeline {
                     steps {
                         archiveArtifacts 'grid-stability-dense.h5'
                         archiveArtifacts 'my_runs/**'
+                        archiveArtifacts 'grid-stability-dense/**'
                     }
                 }
             }

lab5/train/train.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 import mlflow
+from mlflow.models.signature import infer_signature
 import pandas as pd
 from sacred import Experiment
 from sacred.observers import FileStorageObserver, MongoObserver
@@ -47,6 +48,10 @@ def prepare_train_model(learning_rate, _run):
         history = model.fit(X_train, Y_train, epochs=2, validation_data=(X_valid, Y_valid))
 
         model.save('grid-stability-dense.h5')
+        mlflow.keras.save_model(keras_model=model,
+                                path='grid-stability-dense',
+                                signature=infer_signature(X_train, Y_train),
+                                input_example=X_train.iloc[0])
 
         _run.info['history'] = str(history.history['loss'][-1])
         mlflow.log_metric('loss', history.history['loss'][-1])