Update train.py

rishikksh20 · Mar 27, 2021 · d41cb23 · d41cb23
1 parent a458359
commit d41cb23
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -199,8 +199,9 @@ def train(local_rank, args, hp, model):
             if (step + 1) % hp.train.accum_grad == 0:
                 losses.update(loss.item()*hp.train.accum_grad)
                 torch.nn.utils.clip_grad_norm_(model.parameters(), hp.train.grad_clip)
-                scheduler.step()
+
                 optimizer.step()
+                scheduler.step()
                 optimizer.zero_grad()
                 global_step += 1
 
@@ -275,4 +276,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()