For single device move in pre_dispatch after setup function

Author: SeanNaren

pytorch_lightning/plugins/training_type/single_device.py

@@ -64,9 +64,11 @@ def model_to_device(self) -> None:
 
         self._model.to(self.root_device)
 
+    def pre_dispatch(self) -> None:
+        self.model_to_device()
+
     def connect(self, model: torch.nn.Module) -> torch.nn.Module:
         self._model = model
-        self.model_to_device()
         return self.model
 
     @property

pytorch_lightning/trainer/trainer.py

@@ -430,12 +430,8 @@ def fit(
         # ----------------------------
         self.call_hook("on_before_accelerator_backend_setup", model)
         self.accelerator.setup(self, model)  # note: this sets up self.lightning_module
-        # TODO If we call the setup hook here, we might not move our model to the correct device since setup is called
-        # TODO On One GPU
         self.call_setup_hook(model)
 
-
-
         # ----------------------------
         # INSPECT THE CORE LOOPS
         # ----------------------------