Add MAE loss test.

PiperOrigin-RevId: 574830522
google · Oct 19, 2023 · ba12f1b · ba12f1b
1 parent 9149829
commit ba12f1b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,7 +4,7 @@
 
 ### Feature
 
--   Mean average error loss for GBT.
+-   Add support for Mean average error (MAE) loss for GBT.
 
 ## 1.6.0 - 2023-09-28
 

diff --git a/yggdrasil_decision_forests/serving/decision_forest/quick_scorer_extended.cc b/yggdrasil_decision_forests/serving/decision_forest/quick_scorer_extended.cc
@@ -933,7 +933,8 @@ template <>
 absl::Status GenericToSpecializedModel(
     const model::gradient_boosted_trees::GradientBoostedTreesModel& src,
     GradientBoostedTreesRegressionQuickScorerExtended* dst) {
-  if (src.loss() != Loss::SQUARED_ERROR) {
+  if (src.loss() != Loss::SQUARED_ERROR &&
+      src.loss() != Loss::MEAN_AVERAGE_ERROR) {
     return absl::InvalidArgumentError(
         "The GBDT is not trained for regression with squared error loss.");
   }