Update lr_schedulers.py

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