diff --git a/samples/model-builder/create_training_pipeline_tabular_classification_sample.py b/samples/model-builder/create_training_pipeline_tabular_classification_sample.py
index 6bd9405383..317a1472b6 100644
--- a/samples/model-builder/create_training_pipeline_tabular_classification_sample.py
+++ b/samples/model-builder/create_training_pipeline_tabular_classification_sample.py
@@ -33,6 +33,7 @@ def create_training_pipeline_tabular_classification_sample(
 
     tabular_classification_job = aiplatform.AutoMLTabularTrainingJob(
         display_name=display_name,
+        optimization_prediction_type="classification"
     )
 
     my_tabular_dataset = aiplatform.TabularDataset(dataset_id)
diff --git a/samples/model-builder/create_training_pipeline_tabular_classification_sample_test.py b/samples/model-builder/create_training_pipeline_tabular_classification_sample_test.py
index c015e99785..1cfbb02202 100644
--- a/samples/model-builder/create_training_pipeline_tabular_classification_sample_test.py
+++ b/samples/model-builder/create_training_pipeline_tabular_classification_sample_test.py
@@ -44,6 +44,7 @@ def test_create_training_pipeline_tabular_classification_sample(
     )
     mock_get_automl_tabular_training_job.assert_called_once_with(
         display_name=constants.DISPLAY_NAME,
+        optimization_prediction_type="classification"
     )
     mock_run_automl_tabular_training_job.assert_called_once_with(
         dataset=mock_tabular_dataset,