diff --git a/Jenkinsfile-train b/Jenkinsfile-train
index de37f9d..8995392 100644
--- a/Jenkinsfile-train
+++ b/Jenkinsfile-train
@@ -41,7 +41,7 @@ pipeline {
         stage('Archive model and evaluate it') {
             steps {
                archiveArtifacts artifacts: 'model_out', onlyIfSuccessful: true
-               archiveArtifacts artifacts: 'sacred_runs/**', onlyIfSuccessful: true
+               archiveArtifacts artifacts: 'mlruns/**', onlyIfSuccessful: true
                build job: "s444409-evaluation/${params.BRANCH}/", parameters: [string(name: 'BRANCH', value: "${params.BRANCH}")]
             }
         }
diff --git a/train_model.py b/train_model.py
index 5bdd938..d3af166 100644
--- a/train_model.py
+++ b/train_model.py
@@ -1,6 +1,7 @@
 import torch
 import mlflow
 import argparse
+import numpy as np
 from torch import nn
 from torch.utils.data import DataLoader
 from urllib.parse import urlparse
@@ -34,6 +35,8 @@ if __name__ == "__main__":
     plant_test = PlantsDataset('data/Plant_1_Generation_Data.csv.test')
     plant_train = PlantsDataset('data/Plant_1_Generation_Data.csv.train')
 
+    input_example = np.array([plant_test.x_train.numpy()[0]])
+
     train_dataloader = DataLoader(plant_train, batch_size=batch_size)
     test_dataloader = DataLoader(plant_test, batch_size=batch_size)
 
@@ -71,9 +74,10 @@ if __name__ == "__main__":
         if tracking_url_type_store != "file":
             mlflow.pytorch.log_model(
                 model,
-                "s444409-power-plant-model",
-                registered_model_name="s444409PowerPlant",
-                signature=signature
+                "s444409",
+                registered_model_name="s444409",
+                signature=signature,
+                input_example=input_example
             )
         else:
-            mlflow.pytorch.log_model(model, "s444409-power-plant-model", signature=signature)
+            mlflow.pytorch.log_model(model, "s444409", signature=signature, input_example=input_example)