Add additional model to device hook if any additional parameters have been set

Author: SeanNaren

pytorch_lightning/plugins/training_type/single_device.py

@@ -65,10 +65,12 @@ def model_to_device(self) -> None:
         self._model.to(self.root_device)
 
     def pre_dispatch(self) -> None:
+        # Ensures any additional parameters defined in setup are moved to the correct device.
         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