fix NDCG scatter_reduce_ arg (pytorch#2547)

Summary: Pull Request resolved: pytorch#2547 tsia https://pytorch.org/docs/stable/generated/torch.Tensor.scatter_reduce_.html#torch.Tensor.scatter_reduce_ Reviewed By: wilson100hong Differential Revision: D65696513 fbshipit-source-id: fb3f56f599959697ae9ab9bee21c6268ff8aa4ea
iamzainhuda · Nov 12, 2024 · be4b9c7 · be4b9c7
1 parent a97cf28
commit be4b9c7
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",
         )
     )