Skip to content

Commit ca59e4e

Browse files
authored
Configure optimizer typing fix follow-up (#20486)
1 parent 64d0d02 commit ca59e4e

File tree

2 files changed

+9
-7
lines changed

2 files changed

+9
-7
lines changed

docs/source-pytorch/conf.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -462,7 +462,9 @@ def _load_py_module(name: str, location: str) -> ModuleType:
462462
("py:obj", "lightning.pytorch.utilities.memory.is_out_of_cpu_memory"),
463463
("py:func", "lightning.pytorch.utilities.rank_zero.rank_zero_only"),
464464
("py:class", "lightning.pytorch.utilities.types.LRSchedulerConfig"),
465-
("py:class", "lightning.pytorch.utilities.types.OptimizerLRSchedulerConfig"),
465+
("py:class", "lightning.pytorch.utilities.types.LRSchedulerConfigType"),
466+
("py:class", "lightning.pytorch.utilities.types.OptimizerConfigType"),
467+
("py:class", "lightning.pytorch.utilities.types.OptimizerLRSchedulerConfigType"),
466468
("py:class", "lightning_habana.pytorch.plugins.precision.HPUPrecisionPlugin"),
467469
("py:class", "lightning_habana.pytorch.strategies.HPUDDPStrategy"),
468470
("py:class", "lightning_habana.pytorch.strategies.HPUParallelStrategy"),

src/lightning/pytorch/utilities/types.py

+6-6
Original file line numberDiff line numberDiff line change
@@ -104,11 +104,11 @@ class LRSchedulerConfigType(TypedDict, total=False):
104104
strict: bool
105105

106106

107-
class OptimizerConfigDict(TypedDict):
107+
class OptimizerConfigType(TypedDict):
108108
optimizer: Optimizer
109109

110110

111-
class OptimizerLRSchedulerConfigDict(TypedDict):
111+
class OptimizerLRSchedulerConfigType(TypedDict):
112112
optimizer: Optimizer
113113
lr_scheduler: Union[LRSchedulerTypeUnion, LRSchedulerConfigType]
114114
monitor: NotRequired[str]
@@ -119,10 +119,10 @@ class OptimizerLRSchedulerConfigDict(TypedDict):
119119
Optimizer,
120120
Sequence[Optimizer],
121121
tuple[Sequence[Optimizer], Sequence[Union[LRSchedulerTypeUnion, LRSchedulerConfig]]],
122-
OptimizerConfigDict,
123-
OptimizerLRSchedulerConfigDict,
124-
Sequence[OptimizerConfigDict],
125-
Sequence[OptimizerLRSchedulerConfigDict],
122+
OptimizerConfigType,
123+
OptimizerLRSchedulerConfigType,
124+
Sequence[OptimizerConfigType],
125+
Sequence[OptimizerLRSchedulerConfigType],
126126
]
127127
]
128128

0 commit comments

Comments
 (0)