[NFC] polish colossalai/communication/ring.py code style (hpcaitech#630)

edenbuaa · Apr 6, 2022 · 2e11853 · 2e11853
1 parent 01cc941
commit 2e11853
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/colossalai/communication/ring.py b/colossalai/communication/ring.py
@@ -34,15 +34,13 @@ def ring_forward(tensor_send_next: torch.Tensor, parallel_mode: ParallelMode):
                                    dtype=tensor_send_next.dtype)
 
     # send to next rank
-    send_next_op = torch.distributed.P2POp(
-        torch.distributed.isend, tensor_send_next,
-        gpc.get_next_global_rank(parallel_mode))
+    send_next_op = torch.distributed.P2POp(torch.distributed.isend, tensor_send_next,
+                                           gpc.get_next_global_rank(parallel_mode))
     ops.append(send_next_op)
 
     # receive from prev rank
-    recv_prev_op = torch.distributed.P2POp(
-        torch.distributed.irecv, tensor_recv_prev,
-        gpc.get_prev_global_rank(parallel_mode))
+    recv_prev_op = torch.distributed.P2POp(torch.distributed.irecv, tensor_recv_prev,
+                                           gpc.get_prev_global_rank(parallel_mode))
     ops.append(recv_prev_op)
 
     if current_rank % 2 == 0: