Fixed weight column parameter name

googleapis · Nov 17, 2020 · 36d6b7c · 36d6b7c
1 parent e709025
commit 36d6b7c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git 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
@@ -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,