fix: RET505

Author: zhiyuan1i

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):

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):