add return type

Lightning-AI · Feb 21, 2021 · 65b0fe2 · 65b0fe2
1 parent ce6b96e
commit 65b0fe2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
@@ -38,6 +38,7 @@
 from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.core.saving import ALLOWED_CONFIG_TYPES, ModelIO, PRIMITIVE_TYPES
 from pytorch_lightning.core.step_result import Result
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.apply_func import apply_to_collection, convert_to_tensors
 from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
@@ -169,7 +170,7 @@ def automatic_optimization(self) -> bool:
         return self._automatic_optimization
 
     @property
-    def running_stage(self):
+    def running_stage(self) -> Optional[RunningStage]:
         return self.trainer._running_stage if self.trainer else None
 
     @automatic_optimization.setter