fix: RET505

Lightning-AI · Dec 20, 2024 · 4dd9e25 · 4dd9e25
1 parent dc6430c
commit 4dd9e25
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/src/lightning/fabric/strategies/ddp.py b/src/lightning/fabric/strategies/ddp.py
@@ -242,11 +242,11 @@ def _create_stream_context(self, device_ids=None):
             and device_ids is not None
         ):
             stream = self.torch_lib.Stream()
-            return self.torch_lib.stream(stream)
+            ctx = self.torch_lib.stream(stream)
         else:
             from contextlib import nullcontext
-
-            return nullcontext()
+            ctx = nullcontext()
+        return ctx
 
 
 class _DDPBackwardSyncControl(_BackwardSyncControl):

diff --git a/src/lightning/pytorch/strategies/ddp.py b/src/lightning/pytorch/strategies/ddp.py
@@ -431,11 +431,11 @@ def _create_stream_context(self, device_ids=None):
             and device_ids is not None
         ):
             stream = self.torch_lib.Stream()
-            return self.torch_lib.stream(stream)
+            ctx = self.torch_lib.stream(stream)
         else:
             from contextlib import nullcontext
-
-            return nullcontext()
+            ctx = nullcontext()
+        return ctx
 
 
 class _DDPForwardRedirection(_ForwardRedirection):