diff --git a/src/lr_schedulers.py b/src/lr_schedulers.py
index 894fbee..583ba4c 100644
--- a/src/lr_schedulers.py
+++ b/src/lr_schedulers.py
@@ -6,9 +6,9 @@
 def init_lr_scheduler(
     optimizer,
     lr_scheduler,  # learning rate scheduler
+    epochs,
     stepsize=[20, 40],  # step size to decay learning rate
     gamma=0.1,  # learning rate decay
-    epochs=600,
 ):
     '''
     Read this line of code: https://github.com/pytorch/vision/blob/main/references/classification/train.py#L304