diff --git a/predict.py b/predict.py
index c4c2444..ff2c2c5 100644
--- a/predict.py
+++ b/predict.py
@@ -17,9 +17,9 @@ predictions = model.predict(X_test)
 with open("predictions.txt", "w") as f:
     f.write(str(predictions))
 
-accuracy = root_mean_squared_error(y_test, predictions)
+rmse = root_mean_squared_error(y_test, predictions)
 with open("rmse.txt", 'a') as file:
-    file.write(str(accuracy)+"\n")
+    file.write(str(rmse)+"\n")
 
 with open("rmse.txt", 'r') as file:
     lines = file.readlines()
diff --git a/requirements.txt b/requirements.txt
index 466028b..8ae3aeb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,4 +2,5 @@ pandas
 scikit-learn
 tensorflow
 numpy
-matplotlib
\ No newline at end of file
+matplotlib
+mlflow
\ No newline at end of file
diff --git a/train.py b/train.py
index 2f8b27d..babcd5c 100644
--- a/train.py
+++ b/train.py
@@ -2,7 +2,9 @@ import pandas as pd
 from tensorflow import keras
 from tensorflow.keras import layers
 import argparse
-
+import mlflow
+import mlflow.sklearn
+mlflow.set_experiment("s464980")
 class RegressionModel:
     def __init__(self, optimizer="adam", loss="mean_squared_error"):
         self.model = keras.Sequential([
@@ -26,7 +28,6 @@ class RegressionModel:
         self.y_test = data_test["Performance Index"]
 
     def train(self, epochs=30):
-
         self.model.compile(optimizer=self.optimizer, loss=self.loss)
         self.model.fit(self.X_train, self.y_train, epochs=epochs, batch_size=32, validation_data=(self.X_test, self.y_test))
 
@@ -37,6 +38,7 @@ class RegressionModel:
     def evaluate(self):
         test_loss = self.model.evaluate(self.X_test, self.y_test)
         print(f"Test Loss: {test_loss:.4f}")
+        return test_loss
 
     def save_model(self):
         self.model.save("model.keras")
@@ -48,6 +50,9 @@ parser.add_argument('--epochs')
 args = parser.parse_args()
 model = RegressionModel()
 model.load_data("df_train.csv", "df_test.csv")
-model.train(epochs=int(args.epochs))
-model.evaluate()
+with mlflow.start_run() as run:
+    model.train(epochs=int(args.epochs))
+    rmse = model.evaluate()
+    mlflow.log_param("epoch", int(args.epochs))
+    mlflow.log_metric("rmse", rmse)
 model.save_model()