Bug fix for local run in TensorAllToAllValuesAwaitable (#2211)

Summary:
Pull Request resolved: #2211

bug fix when local workers is 1. we should take in variable ids instead of one with splits

Reviewed By: sarckk

Differential Revision: D59491275

fbshipit-source-id: dbe290b4ec250d8e04aa5dc93d5096d6c829d3ca

torchrec/distributed/dist_data.py

@@ -1558,7 +1558,7 @@ def __init__(
 
         self._dist_values: torch.Tensor
         if self._workers == 1:
-            self._dist_values = input_splits
+            self._dist_values = input
             return
         else:
             if input.dim() > 1: