move state dict to cpu before converting to state stream (#208)

Co-authored-by: mspanring <[email protected]>

Author: MarkusSpanring

ray_lightning/launchers/ray_launcher.py

@@ -301,10 +301,9 @@ def _wrapping_function(
         results = function(*args, **kwargs)
 
         if trainer is not None:
-            results = self._collect_rank_zero_results(trainer, results)
-
-        if trainer.strategy.local_rank == 0:
-            return move_data_to_device(results, "cpu")
+            return self._collect_rank_zero_results(trainer, results)
+        else:
+            return None
 
         trainer._teardown()
         trainer._call_teardown_hook()
@@ -326,6 +325,10 @@ def _collect_rank_zero_results(self, trainer: "pl.Trainer",
         if self._strategy.global_rank != 0:
             return None
 
+        # Move state_dict to cpu before converting it to model state stream
+        if trainer.strategy.local_rank == 0:
+            state_dict = move_data_to_device(state_dict, "cpu")
+
         # PyTorch Lightning saves the model weights in a temp file and
         # loads it back on the driver.
         # This won't work in a multi-node setup though, so we return the