Fix typing

cybershiptrooper · Sep 13, 2024 · e6d278c · e6d278c
1 parent 7512a7e
commit e6d278c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/iit/model_pairs/base_model_pair.py b/iit/model_pairs/base_model_pair.py
@@ -273,7 +273,7 @@ def train(
                 del scheduler_kwargs['final_lr']
 
             if "lr_lambda" not in scheduler_kwargs:
-                def linear_lr(step):
+                def linear_lr(step: int) -> float:
                     return 1 - (step / epochs) * (1 - final_lr / initial_lr)
 
                 scheduler_kwargs["lr_lambda"] = linear_lr