make sure that only inference params are allowed, not train ones

Galileo-Galilei · Nov 23, 2024 · 25010aa · 25010aa
1 parent cb7e9f5
commit 25010aa
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/kedro_mlflow/framework/hooks/mlflow_hook.py b/kedro_mlflow/framework/hooks/mlflow_hook.py
@@ -394,7 +394,7 @@ def after_pipeline_run(
                             # I add the special "runner" parameter to be able to choose it at runtime
                             pipeline_params = {
                                 ds_name[7:]: catalog.load(ds_name)
-                                for ds_name in pipeline.inputs()
+                                for ds_name in pipeline.inference.inputs()
                                 if ds_name.startswith("params:")
                             } | {"runner": "SequentialRunner"}
                             model_signature = infer_signature(