update ddp

Lightning-AI · Dec 20, 2024 · ee8935e · ee8935e
1 parent 396ce2c
commit ee8935e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 11 deletions.
diff --git a/src/lightning/fabric/strategies/ddp.py b/src/lightning/fabric/strategies/ddp.py
@@ -233,12 +233,7 @@ def _create_stream_context(self, device_ids=None):
 
         torch_lib = getattr(torch, self.root_device.type)
         # Check if the device type supports streams and has the necessary attributes.
-        if (
-            hasattr(torch_lib, "Stream")
-            and hasattr(torch_lib, "stream")
-            and device_ids is not None
-        ):
-            # ctx = torch.cuda.stream(torch.cuda.Stream()) if device_ids is not None else nullcontext()
+        if hasattr(torch_lib, "Stream") and hasattr(torch_lib, "stream") and device_ids is not None:
             stream = torch_lib.Stream()
             ctx = torch_lib.stream(stream)
         else:

diff --git a/src/lightning/pytorch/strategies/ddp.py b/src/lightning/pytorch/strategies/ddp.py
@@ -423,11 +423,7 @@ def _create_stream_context(self, device_ids=None):
 
         torch_lib = getattr(torch, self.root_device.type)
         # Check if the device type supports streams and has the necessary attributes.
-        if (
-            hasattr(torch_lib, "Stream")
-            and hasattr(torch_lib, "stream")
-            and device_ids is not None
-        ):
+        if hasattr(torch_lib, "Stream") and hasattr(torch_lib, "stream") and device_ids is not None:
             # ctx = torch.cuda.stream(torch.cuda.Stream()) if device_ids is not None else nullcontext()
             stream = torch_lib.Stream()
             ctx = torch_lib.stream(stream)