Skip to content

Commit 9d62118

Browse files
haru-256rwightman
authored andcommitted
refactor(scheduler): add type hints to CosineLRScheduler
- Add type annotations to `__init__` arguments in `CosineLRScheduler` - Import `Tuple` and `Union` from `typing` - Fix missing newline at end of file
1 parent d4ab516 commit 9d62118

File tree

1 file changed

+12
-12
lines changed

1 file changed

+12
-12
lines changed

timm/scheduler/cosine_lr.py

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import math
99
import numpy as np
1010
import torch
11-
from typing import List
11+
from typing import Tuple, List, Union
1212

1313
from .scheduler import Scheduler
1414

@@ -35,16 +35,16 @@ def __init__(
3535
cycle_mul: float = 1.,
3636
cycle_decay: float = 1.,
3737
cycle_limit: int = 1,
38-
warmup_t=0,
39-
warmup_lr_init=0,
40-
warmup_prefix=False,
41-
t_in_epochs=True,
42-
noise_range_t=None,
43-
noise_pct=0.67,
44-
noise_std=1.0,
45-
noise_seed=42,
46-
k_decay=1.0,
47-
initialize=True,
38+
warmup_t: int=0,
39+
warmup_lr_init: float=0.,
40+
warmup_prefix: bool=False,
41+
t_in_epochs: bool=True,
42+
noise_range_t: Union[List[int], Tuple[int, int], int, None]=None,
43+
noise_pct: float=0.67,
44+
noise_std: float=1.0,
45+
noise_seed: int=42,
46+
k_decay: float=1.0,
47+
initialize: bool=True,
4848
) -> None:
4949
super().__init__(
5050
optimizer,
@@ -114,4 +114,4 @@ def get_cycle_length(self, cycles=0):
114114
t = self.t_initial * cycles
115115
else:
116116
t = int(math.floor(-self.t_initial * (self.cycle_mul ** cycles - 1) / (1 - self.cycle_mul)))
117-
return t + self.warmup_t if self.warmup_prefix else t
117+
return t + self.warmup_t if self.warmup_prefix else t

0 commit comments

Comments
 (0)