fix mypy typing errors in pytorch_lightning/strategies/single_device.py (#13532)

* fix typing in strategies/single_device.py
* Make assert statement more explicit

Co-authored-by: Justus Schock <[email protected]>

Author: CyprienRicque

pyproject.toml

@@ -74,7 +74,6 @@ module = [
     "pytorch_lightning.strategies.parallel",
     "pytorch_lightning.strategies.sharded",
     "pytorch_lightning.strategies.sharded_spawn",
-    "pytorch_lightning.strategies.single_device",
     "pytorch_lightning.strategies.single_tpu",
     "pytorch_lightning.strategies.tpu_spawn",
     "pytorch_lightning.strategies.strategy",

src/pytorch_lightning/strategies/single_device.py

@@ -21,7 +21,7 @@
 import pytorch_lightning as pl
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
-from pytorch_lightning.strategies.strategy import Strategy
+from pytorch_lightning.strategies.strategy import Strategy, TBroadcast
 from pytorch_lightning.utilities.types import _DEVICE
 
 
@@ -66,6 +66,7 @@ def root_device(self) -> torch.device:
         return self._root_device
 
     def model_to_device(self) -> None:
+        assert self.model is not None, "self.model must be set before self.model.to()"
         self.model.to(self.root_device)
 
     def setup(self, trainer: pl.Trainer) -> None:
@@ -76,10 +77,10 @@ def setup(self, trainer: pl.Trainer) -> None:
     def is_global_zero(self) -> bool:
         return True
 
-    def barrier(self, *args, **kwargs) -> None:
+    def barrier(self, *args: Any, **kwargs: Any) -> None:
         pass
 
-    def broadcast(self, obj: object, src: int = 0) -> object:
+    def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
         return obj
 
     @classmethod