add mlflow

train.py

@@ -3,8 +3,8 @@ 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([
@@ -48,8 +48,8 @@ parser = argparse.ArgumentParser()
 parser.add_argument('--epochs')
 
 args = parser.parse_args()
+mlflow.set_tracking_uri("http://localhost:5000")
 model = RegressionModel()
-model.load_data("df_train.csv", "df_test.csv")
 with mlflow.start_run() as run:
     model.train(epochs=int(args.epochs))
     rmse = model.evaluate()