Put TensorPool OutputDist to Remote RO to avoid wrong rebatching on indexing (#3087)

Summary: Pull Request resolved: #3087

Reviewed By: yumin829928

Differential Revision: D76324557

fbshipit-source-id: 1ef856b4c9d626d189b30cef9acd5071142cb1b4

Author: tissue3

torchrec/distributed/tensor_pool.py

@@ -35,6 +35,15 @@
 from torchrec.modules.utils import deterministic_dedup
 
 
+@torch.fx.wrap
+def index_select_view(
+    output: torch.Tensor,
+    unbucketize_permute: Optional[torch.Tensor],
+    dim: int,
+) -> torch.Tensor:
+    return output[unbucketize_permute].view(-1, dim)
+
+
 class TensorPoolAwaitable(LazyAwaitable[torch.Tensor]):
     def __init__(
         self,
@@ -441,7 +450,7 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor:
 
         output = self._lookup_values_dist(lookup_list)
 
-        return output[unbucketize_permute].view(-1, self._dim)
+        return index_select_view(output, unbucketize_permute, self._dim)
 
     # pyre-ignore
     def _update_values_dist(self, ctx: ObjectPoolShardingContext, values: torch.Tensor):