From 2d53175cd692f97e4a5f3fdc34a15834f9216468 Mon Sep 17 00:00:00 2001
From: Zain Huda <zainhuda@meta.com>
Date: Fri, 8 Nov 2024 16:27:48 -0800
Subject: [PATCH] fix NDCG scatter_reduce_ arg

Summary: tsia

Differential Revision: D65696513
---
 torchrec/metrics/ndcg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchrec/metrics/ndcg.py b/torchrec/metrics/ndcg.py
index 816a8e678..61be6525d 100644
--- 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",
         )
     )