chore: simplify try-finally logic in pynccl

PygmalionAI · Jul 1, 2024 · d1a3c7b · d1a3c7b
1 parent 440384d
commit d1a3c7b
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/aphrodite/distributed/device_communicators/pynccl.py b/aphrodite/distributed/device_communicators/pynccl.py
@@ -247,15 +247,11 @@ def __init__(
         # now the `device` object is a `torch.device` object
         assert isinstance(device, torch.device)
         self.device = device
-        current_device = torch.cuda.current_device()
-        try:
-            torch.cuda.set_device(self.device)
+        with torch.cuda.device(device):
             NCCL_CHECK(
                 _c_ncclCommInitRank(ctypes.byref(self.comm), self.world_size,
                                     self.unique_id, self.rank))
             self.stream = torch.cuda.Stream()
-        finally:
-            torch.cuda.set_device(current_device)
 
     def all_reduce(self,
                    tensor: torch.Tensor,