fix NDCG scatter_reduce_ arg (#2547)

iamzainhuda · facebook-github-bot · commit be4b9c743128 · 2024-11-12T10:35:36.000-08:00
Summary: Pull Request resolved: #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
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",
         )
     )
 

Original file line number	Diff line number	Diff line change
`@@ -253,7 +253,7 @@ def _get_ndcg_states(`
`253`	`253`	`dim=-1,`
`254`	`254`	`index=expanded_session_ids,`
`255`	`255`	`src=adjusted_weights, # [num_tasks, batch_size]`
`256`		`- reduce="max",`
	`256`	`+ reduce="amax",`
`257`	`257`	`)`
`258`	`258`	`)`
`259`	`259`