diff --git a/cli/jobs/pipelines-with-components/pipeline_with_hyperparameter_sweep/predict.yml b/cli/jobs/pipelines-with-components/pipeline_with_hyperparameter_sweep/predict.yml
index 8bce36d527..a49baa9e02 100644
--- a/cli/jobs/pipelines-with-components/pipeline_with_hyperparameter_sweep/predict.yml
+++ b/cli/jobs/pipelines-with-components/pipeline_with_hyperparameter_sweep/predict.yml
@@ -6,7 +6,7 @@ version: 1
 type: command
 inputs:
   model: 
-    type: mlflow_model
+    type: uri_folder
   test_data:
     type: uri_folder
 outputs:
diff --git a/cli/jobs/pipelines-with-components/pipeline_with_hyperparameter_sweep/train.yml b/cli/jobs/pipelines-with-components/pipeline_with_hyperparameter_sweep/train.yml
index f3187b763e..b726edaca3 100644
--- a/cli/jobs/pipelines-with-components/pipeline_with_hyperparameter_sweep/train.yml
+++ b/cli/jobs/pipelines-with-components/pipeline_with_hyperparameter_sweep/train.yml
@@ -53,7 +53,7 @@ inputs:
 
 outputs:
   model_output:
-    type: mlflow_model
+    type: uri_folder
   test_data:
     type: uri_folder