Fix train_pipeline test (#2299)

Summary:
Pull Request resolved: #2299

Test is flaky: https://www.internalfb.com/intern/test/844425089856635?ref_report_id=0. Order of preproc modules is not always set even though test had it set. This removes the nondeterminism from the test case

Reviewed By: sarckk

Differential Revision: D61299548

fbshipit-source-id: 46475f81077769000fd731c9e23fc7c84e3fd54f

torchrec/distributed/train_pipeline/tests/test_train_pipelines.py

@@ -973,10 +973,16 @@ def test_pipeline_preproc_not_shared_with_arg_transform(self) -> None:
 
         # preproc args
         self.assertEqual(len(pipeline._pipelined_preprocs), 2)
-        for i, input_attr_name in [(0, "idlist_features"), (1, "idscore_features")]:
+        input_attr_names = {"idlist_features", "idscore_features"}
+        for i in range(len(pipeline._pipelined_preprocs)):
             preproc_mod = pipeline._pipelined_preprocs[i]
             self.assertEqual(len(preproc_mod._args), 1)
+
+            input_attr_name = preproc_mod._args[0].input_attrs[1]
+            self.assertTrue(input_attr_name in input_attr_names)
             self.assertEqual(preproc_mod._args[0].input_attrs, ["", input_attr_name])
+            input_attr_names.remove(input_attr_name)
+
             self.assertEqual(preproc_mod._args[0].is_getitems, [False, False])
             # no parent preproc module in FX graph
             self.assertEqual(preproc_mod._args[0].preproc_modules, [None, None])