Fix shape issue with remapped indices in sharded ManagedCollisionEmbeddingCollection (#3010)

Summary:
Pull Request resolved: #3010

Fixing shapes of remapped indices returned by ManagedCollisionEmbeddingCollection to be [B, 1] instead of [B], which causes an issue in SequenceEmbeddingsAllToAll (embedding_dim=a2a_sequence_embs_tensor.shape[1]) related to the tuple index being out of range.

Reviewed By: cx-yin, kausv, xing-liu

Differential Revision:
D75251224

Privacy Context Container: L1292699

fbshipit-source-id: 5afa863f778a74f6fbb51a9afba71986eff8c7e2

Author: Arthi Suresh

torchrec/distributed/mc_modules.py

@@ -671,7 +671,9 @@ def _kjt_list_to_tensor_list(
                     vals.append(feature_split.values() + offset)
                 remapped_ids_ret.append(torch.cat(vals).view(-1, 1))
             else:
-                remapped_ids_ret.append(kjt.values() + self._table_to_offset[tables[0]])
+                remapped_ids_ret.append(
+                    (kjt.values() + self._table_to_offset[tables[0]]).unsqueeze(-1)
+                )
         return remapped_ids_ret
 
     def global_to_local_index(