bug fix

Author: Bissmella

src/diffusers/models/attention_dispatch.py

@@ -1036,7 +1036,7 @@ def _all_to_all_dim_exchange(x: torch.Tensor, scatter_idx: int = 2, gather_idx:
 
         if group_world_size >1:
             #maybe here need to use the _all_to_all_single helper to avoid contiguity issues
-            out = _all_to_all_single(x_temp, None, None, group=group)
+            out = _all_to_all_single(x_temp, group=group)
             #out = _wait_tensor(out)
         else:
             out = x_temp
@@ -1055,7 +1055,7 @@ def _all_to_all_dim_exchange(x: torch.Tensor, scatter_idx: int = 2, gather_idx:
 
         if group_world_size >1:
             #maybe here need to use the _all_to_all_single helper to avoid contiguity issues
-            output = _all_to_all_single(x_temp, None, None, group)
+            output = _all_to_all_single(x_temp, group)
             #output = _wait_tensor(output)
         else:
             output = x_temp