1 files changed, 5 insertions, 5 deletions

diff --git a/training/optimization.py b/training/optimization.py
index 5db7794..6dee4bc 100644
--- a/training/optimization.py
+++ b/training/optimization.py
@@ -49,8 +49,8 @@ def get_one_cycle_schedule(
     annealing: Literal["cos", "half_cos", "linear"] = "cos",
     warmup_exp: int = 1,
     annealing_exp: int = 1,
-    min_lr: int = 0.04,
-    mid_point: int = 0.3,
+    min_lr: float = 0.04,
+    mid_point: float = 0.3,
     last_epoch: int = -1
 ):
     if warmup == "linear":
@@ -91,10 +91,10 @@ def get_scheduler(
     id: str,
     optimizer: torch.optim.Optimizer,
     num_training_steps_per_epoch: int,
-    gradient_accumulation_steps: int,
+    gradient_accumulation_steps: int = 1,
     min_lr: float = 0.04,
-    warmup_func: str = "cos",
-    annealing_func: str = "cos",
+    warmup_func: Literal["cos", "linear"] = "cos",
+    annealing_func: Literal["cos", "half_cos", "linear"] = "cos",
     warmup_exp: int = 1,
     annealing_exp: int = 1,
     cycles: int = 1,