From 36d6b7ca1c365deaf320886e353fa7511bcaff5f Mon Sep 17 00:00:00 2001
From: Ivan Cheung <ivanmkc@google.com>
Date: Tue, 17 Nov 2020 19:30:09 +0900
Subject: [PATCH] Fixed weight column parameter name
---
google/cloud/aiplatform/training_jobs.py | 2 +-
tests/unit/aiplatform/test_automl_tables_training_jobs.py | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py
index 65d8f55bc2..067761fb12 100644
--- a/google/cloud/aiplatform/training_jobs.py
+++ b/google/cloud/aiplatform/training_jobs.py
@@ -1108,7 +1108,7 @@ def run(
"transformations": self._column_transformations,
"trainBudgetMilliNodeHours": budget_milli_node_hours,
# optional inputs
- "weightColumn": weight_column,
+ "weightColumnName": weight_column,
"disableEarlyStopping": disable_early_stopping,
"optimizationObjective": self._optimization_objective,
"predictionType": self._optimization_prediction_type,
diff --git a/tests/unit/aiplatform/test_automl_tables_training_jobs.py b/tests/unit/aiplatform/test_automl_tables_training_jobs.py
index 109890447d..211c513185 100644
--- a/tests/unit/aiplatform/test_automl_tables_training_jobs.py
+++ b/tests/unit/aiplatform/test_automl_tables_training_jobs.py
@@ -52,7 +52,7 @@
"transformations": _TEST_TRAINING_COLUMN_TRANSFORMATIONS,
"trainBudgetMilliNodeHours": _TEST_TRAINING_BUDGET_MILLI_NODE_HOURS,
# optional inputs
- "weightColumn": _TEST_TRAINING_WEIGHT_COLUMN,
+ "weightColumnName": _TEST_TRAINING_WEIGHT_COLUMN,
"disableEarlyStopping": _TEST_TRAINING_DISABLE_EARLY_STOPPING,
"predictionType": _TEST_TRAINING_OPTIMIZATION_PREDICTION_TYPE,
"optimizationObjective": _TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME,