diff --git a/samples/model-builder/create_training_pipeline_tabular_regression_sample.py b/samples/model-builder/create_training_pipeline_tabular_regression_sample.py
index 2404bb37e2..f7edcce1d9 100644
--- a/samples/model-builder/create_training_pipeline_tabular_regression_sample.py
+++ b/samples/model-builder/create_training_pipeline_tabular_regression_sample.py
@@ -33,6 +33,7 @@ def create_training_pipeline_tabular_regression_sample(
 
     tabular_regression_job = aiplatform.AutoMLTabularTrainingJob(
         display_name=display_name,
+        optimization_prediction_type="regression"
     )
 
     my_tabular_dataset = aiplatform.TabularDataset(dataset_id)
diff --git a/samples/model-builder/create_training_pipeline_tabular_regression_sample_test.py b/samples/model-builder/create_training_pipeline_tabular_regression_sample_test.py
index 1e897b5851..d9a6b386e0 100644
--- a/samples/model-builder/create_training_pipeline_tabular_regression_sample_test.py
+++ b/samples/model-builder/create_training_pipeline_tabular_regression_sample_test.py
@@ -44,6 +44,7 @@ def test_create_training_pipeline_tabular_regression_sample(
     )
     mock_get_automl_tabular_training_job.assert_called_once_with(
         display_name=constants.DISPLAY_NAME,
+        optimization_prediction_type="regression"
     )
     mock_run_automl_tabular_training_job.assert_called_once_with(
         dataset=mock_tabular_dataset,