fix NDCG scatter_reduce_ arg

Summary: tsia Differential Revision: D65696513
iamzainhuda · Nov 9, 2024 · 2d53175 · 2d53175
1 parent ead7c90
commit 2d53175
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrec/metrics/ndcg.py b/torchrec/metrics/ndcg.py
@@ -253,7 +253,7 @@ def _get_ndcg_states(
             dim=-1,
             index=expanded_session_ids,
             src=adjusted_weights,  # [num_tasks, batch_size]
-            reduce="max",
+            reduce="amax",
         )
     )