diff --git a/src/lightning/pytorch/utilities/types.py b/src/lightning/pytorch/utilities/types.py index 8fccfa79c976a..a67e21796c9b4 100644 --- a/src/lightning/pytorch/utilities/types.py +++ b/src/lightning/pytorch/utilities/types.py @@ -104,9 +104,14 @@ class LRSchedulerConfigType(TypedDict, total=False): strict: bool -class OptimizerLRSchedulerConfig(TypedDict): +class OptimizerConfigDict(TypedDict): optimizer: Optimizer - lr_scheduler: NotRequired[Union[LRSchedulerTypeUnion, LRSchedulerConfigType]] + + +class OptimizerLRSchedulerConfigDict(TypedDict): + optimizer: Optimizer + lr_scheduler: Union[LRSchedulerTypeUnion, LRSchedulerConfigType] + monitor: NotRequired[str] OptimizerLRScheduler = Optional[ @@ -114,8 +119,10 @@ class OptimizerLRSchedulerConfig(TypedDict): Optimizer, Sequence[Optimizer], tuple[Sequence[Optimizer], Sequence[Union[LRSchedulerTypeUnion, LRSchedulerConfig]]], - OptimizerLRSchedulerConfig, - Sequence[OptimizerLRSchedulerConfig], + OptimizerConfigDict, + OptimizerLRSchedulerConfigDict, + Sequence[OptimizerConfigDict], + Sequence[OptimizerLRSchedulerConfigDict], ] ]