Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add return type

This commit is contained in:
Adrian Wälchli 2021-02-21 02:54:11 +01:00
parent ce6b96e045
commit 65b0fe269c
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pytorch_lightning/core/lightning.py
View File

@ -38,6 +38,7 @@ from pytorch_lightning.core.memory import ModelSummary
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 @@ class LightningModule(
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