refactor(scheduler): add type hints to CosineLRScheduler

timm/scheduler/cosine_lr.py

@@ -8,7 +8,7 @@
 import math
 import numpy as np
 import torch
-from typing import List
+from typing import Tuple, List, Union
 
 from .scheduler import Scheduler
 
@@ -35,16 +35,16 @@ def __init__(
             cycle_mul: float = 1.,
             cycle_decay: float = 1.,
             cycle_limit: int = 1,
-            warmup_t=0,
-            warmup_lr_init=0,
-            warmup_prefix=False,
-            t_in_epochs=True,
-            noise_range_t=None,
-            noise_pct=0.67,
-            noise_std=1.0,
-            noise_seed=42,
-            k_decay=1.0,
-            initialize=True,
+            warmup_t: int=0,
+            warmup_lr_init: float=0.,
+            warmup_prefix: bool=False,
+            t_in_epochs: bool=True,
+            noise_range_t: Union[List[int], Tuple[int, int], int, None]=None,
+            noise_pct: float=0.67,
+            noise_std: float=1.0,
+            noise_seed: int=42,
+            k_decay: float=1.0,
+            initialize: bool=True,
     ) -> None:
         super().__init__(
             optimizer,
@@ -114,4 +114,4 @@ def get_cycle_length(self, cycles=0):
             t = self.t_initial * cycles
         else:
             t = int(math.floor(-self.t_initial * (self.cycle_mul ** cycles - 1) / (1 - self.cycle_mul)))
-        return t + self.warmup_t if self.warmup_prefix else t
+        return t + self.warmup_t if self.warmup_prefix else t