From a025fa77d7c8d53e5b1a887fcac0d9549ad76086 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Sun, 3 Jan 2021 21:21:44 +0100
Subject: [PATCH] update schedular after each epoch

---
 pts/trainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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