From de7d9fb18fcc8a43e2d365203904514e89dd414e Mon Sep 17 00:00:00 2001
From: Jules Laplace <julescarbon@gmail.com>
Date: Thu, 6 Sep 2018 11:24:21 +0200
Subject: fix augment

---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'train.py')

diff --git a/train.py b/train.py
index 0fe892d..92ef89d 100755
--- a/train.py
+++ b/train.py
@@ -121,5 +121,5 @@ for epoch in range(start_epoch, start_epoch + opt.niter + opt.niter_decay + 1):
         model.module.update_fixed_params()
 
     ### linearly decay learning rate after certain iterations
-    if epoch > opt.niter:
+    if opt.niter != 0 and epoch > opt.niter:
         model.module.update_learning_rate()
-- 
cgit v1.2.3-70-g09d2