Hello!

Im trying to save my model with mlflow in databricks, it is a xgboost model, when I save it using code it saves with a sklearn flavor and not saves other parameters, also I'm using kedro with kedro-mlflow plugin.

def log_metrics_and_model(model, metrics: pd.DataFrame, params: Dict, X_train: pd.DataFrame):

    model_name = params.get("model_name", "model")
    
    mlflow.log_params(params)
    model_params = model.get_params()
    mlflow.log_params(model_params)

    for metric_name, metric_values in metrics.items():
        mlflow.log_metric(f"{model_name}_{metric_name}_train", metric_values.iloc[0])
        mlflow.log_metric(f"{model_name}_{metric_name}_test", metric_values.iloc[1])
    
    input_example = X_train.iloc[:5]
    signature = infer_signature(X_train, model.predict(X_train.iloc[:5]))

    mlflow.xgboost.log_model(
        model,
        model_name,
        signature=signature,
        input_example=input_example
    )
    

In databricks it saves a simple sklearn model, but locally saves xgboost with all properties correctly, any idea?