Configure optimizer typing fix follow-up (#20486)

Lightning-AI · Dec 10, 2024 · ca59e4e · ca59e4e
1 parent 64d0d02
commit ca59e4e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git 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
@@ -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],
     ]
 ]