Bug fix for local run in TensorAllToAllValuesAwaitable

Summary: bug fix when local workers is 1. we should take in variable ids instead of one with splits Reviewed By: sarckk Differential Revision: D59491275
sarckk · Jul 8, 2024 · b8daf38 · b8daf38
1 parent 7777565
commit b8daf38
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrec/distributed/dist_data.py b/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: