Add additional model to device hook if any additional parameters have…

… been set
Lightning-AI · Mar 17, 2021 · bd2a53a · bd2a53a
1 parent 72097ba
commit bd2a53a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pytorch_lightning/plugins/training_type/single_device.py b/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