[YDF] Rename LAMBDA_MART_NDCG5 to LAMBDA_MART_NDCG

We want to make the truncation parameter configurable. Renaming the loss is the first step. PiperOrigin-RevId: 675162089
tensorflow · Sep 16, 2024 · f253b10 · f253b10
1 parent 16c057f
commit f253b10
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorflow_decision_forests/component/builder/builder.py b/tensorflow_decision_forests/component/builder/builder.py
@@ -949,7 +949,7 @@ def __init__(
 
     elif isinstance(objective, py_tree.objective.RankingObjective):
       num_trees_per_iter = 1
-      loss = gradient_boosted_trees_pb2.Loss.LAMBDA_MART_NDCG5
+      loss = gradient_boosted_trees_pb2.Loss.LAMBDA_MART_NDCG
       bias = [bias]
 
     elif isinstance(objective, py_tree.objective.AbstractUpliftObjective):