diff --git a/pts/trainer.py b/pts/trainer.py
index edf8292..9486253 100644
--- a/pts/trainer.py
+++ b/pts/trainer.py
@@ -91,9 +91,10 @@ class Trainer:
 
                     loss.backward()
                     nn.utils.clip_grad_norm_(net.parameters(), self.clip_gradient)
-                    lr_scheduler.step(loss)
+                    optimizer.step()
 
                     if self.num_batches_per_epoch == batch_no:
+                        lr_scheduler.step(avg_epoch_loss / batch_no)
                         break
 
             # mark epoch end time and log time cost of current epoch