make dict attribute of train input pipelineable (pytorch#2278)

Summary:
Pull Request resolved: pytorch#2278

When the training input has attribute of dict, accessing those with "get" will break the train pipeline. Here we add support to make dict attribute pipelineable.

Reviewed By: ge0405

Differential Revision: D60635872

fbshipit-source-id: f0b2c314ab90a5e4247bc5c75846633b39630cf2

torchrec/distributed/train_pipeline/utils.py

@@ -790,6 +790,12 @@ def _get_node_args_helper(
                     arg = child_node.kwargs["values"]
                 else:
                     arg = child_node.args[1]
+            elif child_node.op == "call_method" and child_node.target == "get":
+                # pyre-ignore[6]
+                arg_info.input_attrs.insert(0, child_node.args[1])
+                arg_info.is_getitems.insert(0, True)
+                arg_info.preproc_modules.insert(0, None)
+                arg = child_node.args[0]
             elif child_node.op == "call_module":
                 preproc_module_fqn = str(child_node.target)
                 preproc_module = getattr(model, preproc_module_fqn, None)