fix: testando correcao de nome de modelo

heltonricardo · Dec 22, 2023 · b0bffcc · b0bffcc
1 parent 17572e1
commit b0bffcc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/train.py b/src/train.py
@@ -151,7 +151,7 @@ def train_model(model, X_train, y_train, is_train=True):
         model.fit(X_train, y_train, epochs=50, validation_split=0.2, verbose=3)
     if is_train:
         run_uri = f"runs:/{run.info.run_id}"
-        mlflow.register_model(run_uri, os.getenv("MLFLOW_TRACKING_MODEL"))
+        mlflow.register_model(run_uri, "fetal_health")
 
 
 if __name__ == "__main__":