Hi guys, Within MLflow I want the model to return (response) both the class and the probability of a request sent. I leave the code below. Thank you. with mlflow.start_run(): dtc = DecisionTreeClassifier(max_depth=1, random_state=10) dtc.fit(X_train, y_train) y_pred_class = dtc.predict(X_test) accuracy = metrics.accuracy_score(y_test, y_pred_class) y_pred_proba = dtc.predict_proba(X_test) aux_df = pd.DataFrame(y_pred_proba) print(f"accuracy: {accuracy}") mlflow.log_param('random_state', 42) mlflow.log_param('max_depth', 10) mlflow.log_param('min_samples_split', 20) mlflow.log_param('min_samples_leaf', 5) mlflow.log_param('criterion', 'gini') mlflow.log_metric('accuracy', accuracy) I suposse is this section: # Register the model with accuracy and probability predictions mlflow.sklearn.log_model(sk_model=dtc, artifact_path='iris_model', pyfunc_predict_fn="predict_proba", await_registration_for=60)

If i may, just copy-pasting ~40 lines of code in an unformatted block, with no kedro-specific question (but rather an ml-flow one?) , will not help the community to assist you. I would suggest you edit (most of) the code away, and leave only the relevant part that is related to kedro, i.e. the part that doesn't work for you. I hope this helps to get a response. 🙃