From ca59e4e1bd7b76a4814db9d183c03513937ecbd0 Mon Sep 17 00:00:00 2001 From: Luca Antiga Date: Tue, 10 Dec 2024 10:46:15 +0100 Subject: [PATCH] Configure optimizer typing fix follow-up (#20486) --- docs/source-pytorch/conf.py | 4 +++- src/lightning/pytorch/utilities/types.py | 12 ++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/docs/source-pytorch/conf.py b/docs/source-pytorch/conf.py index 107f58c5797b6..7f31ea9f4c1c2 100644 --- a/docs/source-pytorch/conf.py +++ b/docs/source-pytorch/conf.py @@ -462,7 +462,9 @@ def _load_py_module(name: str, location: str) -> ModuleType: ("py:obj", "lightning.pytorch.utilities.memory.is_out_of_cpu_memory"), ("py:func", "lightning.pytorch.utilities.rank_zero.rank_zero_only"), ("py:class", "lightning.pytorch.utilities.types.LRSchedulerConfig"), - ("py:class", "lightning.pytorch.utilities.types.OptimizerLRSchedulerConfig"), + ("py:class", "lightning.pytorch.utilities.types.LRSchedulerConfigType"), + ("py:class", "lightning.pytorch.utilities.types.OptimizerConfigType"), + ("py:class", "lightning.pytorch.utilities.types.OptimizerLRSchedulerConfigType"), ("py:class", "lightning_habana.pytorch.plugins.precision.HPUPrecisionPlugin"), ("py:class", "lightning_habana.pytorch.strategies.HPUDDPStrategy"), ("py:class", "lightning_habana.pytorch.strategies.HPUParallelStrategy"), diff --git a/src/lightning/pytorch/utilities/types.py b/src/lightning/pytorch/utilities/types.py index a67e21796c9b4..3fd80d5e07a9a 100644 --- a/src/lightning/pytorch/utilities/types.py +++ b/src/lightning/pytorch/utilities/types.py @@ -104,11 +104,11 @@ class LRSchedulerConfigType(TypedDict, total=False): strict: bool -class OptimizerConfigDict(TypedDict): +class OptimizerConfigType(TypedDict): optimizer: Optimizer -class OptimizerLRSchedulerConfigDict(TypedDict): +class OptimizerLRSchedulerConfigType(TypedDict): optimizer: Optimizer lr_scheduler: Union[LRSchedulerTypeUnion, LRSchedulerConfigType] monitor: NotRequired[str] @@ -119,10 +119,10 @@ class OptimizerLRSchedulerConfigDict(TypedDict): Optimizer, Sequence[Optimizer], tuple[Sequence[Optimizer], Sequence[Union[LRSchedulerTypeUnion, LRSchedulerConfig]]], - OptimizerConfigDict, - OptimizerLRSchedulerConfigDict, - Sequence[OptimizerConfigDict], - Sequence[OptimizerLRSchedulerConfigDict], + OptimizerConfigType, + OptimizerLRSchedulerConfigType, + Sequence[OptimizerConfigType], + Sequence[OptimizerLRSchedulerConfigType], ] ]