From 2005d7faaf0b3a66ff6f7f64d5d0ce957a462b46 Mon Sep 17 00:00:00 2001 From: Zain Huda Date: Tue, 30 Jan 2024 10:56:22 -0800 Subject: [PATCH] dynamically set session_key for session_ids in NDCG (#1658) Summary: Pull Request resolved: https://github.com/pytorch/torchrec/pull/1658 Previously session_key was hardcoded to "session_id" which is unsuitable as all models do not have the same session_key. Now the defined session_key from config will be used as the key instead. Reviewed By: howei Differential Revision: D53067426 fbshipit-source-id: e695b5c882eedcb6737168a8f585fbcdf427d590 --- torchrec/metrics/ndcg.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchrec/metrics/ndcg.py b/torchrec/metrics/ndcg.py index a06c0feac..c77bc029a 100644 --- a/torchrec/metrics/ndcg.py +++ b/torchrec/metrics/ndcg.py @@ -307,4 +307,6 @@ def __init__( process_group=process_group, **kwargs, ) - self._required_inputs.add("session_id") + # session_key is set through front end config + # pyre-ignore[6] + self._required_inputs.add(kwargs["session_key"])