Update train.py

Aman123lug · Aug 28, 2023 · 8134bdc · 8134bdc
1 parent c820065
commit 8134bdc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/train.py b/src/train.py
@@ -55,10 +55,10 @@ def training(self, path:Path):
                 # please refer to the doc for more information:
                 # https://mlflow.org/docs/latest/model-registry.html#api-workflow
                 mlflow.sklearn.log_model(
-                    dtree, "model", registered_model_name="DecisionTree", signature=signature
+                    dtree, "model/model.pkl", registered_model_name="DecisionTree", signature=signature
                 )
             else:
-                mlflow.sklearn.log_model(dtree, "model", signature=signature)
+                mlflow.sklearn.log_model(dtree, "model/model.pkl", signature=signature)
 
         with open("metrics.txt", "w+") as f:
             f.write(str(test_accracy))