Simplify permute indices for Sharded EBC output_dist creation (#2856)

Summary:
Pull Request resolved: #2856

Simplify create_output_dist call See suggestion in next diff

Reviewed By: iamzainhuda

Differential Revision: D72079015

fbshipit-source-id: 37d9bada0b6742db27da69e1b62ab3073145a98e

Author: aporialiao

torchrec/distributed/embeddingbag.py

@@ -1172,16 +1172,13 @@ def _create_output_dist(self) -> None:
         for i, name in enumerate(self._uncombined_embedding_names):
             embedding_name_order.setdefault(name, i)
 
-        def sort_key(input: Tuple[int, str]) -> Tuple[int, int]:
-            index, name = input
-            return (embedding_name_order[name], embedding_shard_offsets[index])
-
-        permute_indices = [
-            i
-            for i, _ in sorted(
-                enumerate(self._uncombined_embedding_names), key=sort_key
-            )
-        ]
+        permute_indices = sorted(
+            range(len(self._uncombined_embedding_names)),
+            key=lambda i: (
+                embedding_name_order[self._uncombined_embedding_names[i]],
+                embedding_shard_offsets[i],
+            ),
+        )
         self._permute_op: PermutePooledEmbeddings = PermutePooledEmbeddings(
             self._uncombined_embedding_dims, permute_indices, self._device
         )