From 6c576630647375393ce14dd2a8c026c61aff6f2f Mon Sep 17 00:00:00 2001 From: Ruilin Chen Date: Wed, 4 Sep 2024 21:55:34 -0700 Subject: [PATCH] Fix ph_key issue with trec pipeline (#2363) Summary: Pull Request resolved: https://github.com/pytorch/torchrec/pull/2363 Fix ph_key issue by: 1. make ph_key for TrainModelInput to be model_input.id_list_features["defalut"] 2. handles parsing in trec utils. Differential Revision: D62160686 --- torchrec/distributed/train_pipeline/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchrec/distributed/train_pipeline/utils.py b/torchrec/distributed/train_pipeline/utils.py index 4e28a12ba..f6c516d35 100644 --- a/torchrec/distributed/train_pipeline/utils.py +++ b/torchrec/distributed/train_pipeline/utils.py @@ -774,7 +774,8 @@ def _get_node_args_helper( # pyre-ignore[16] ph_key: str = child_node.ph_key # example: ph_key = 'event_id_list_features_seqs[marketplace]' - ph_keys = ph_key.split("[") + ph_key = ph_key.replace("[", ".") + ph_keys = ph_key.split(".") for key in ph_keys: if "]" in key: arg_info.input_attrs.append(key[:-1])