train numbers better

1 files changed, 4 insertions, 4 deletions

diff --git a/train.py b/train.py
index dd79bab..8fef794 100644
--- a/train.py
+++ b/train.py
@@ -15,7 +15,7 @@ if __name__ == '__main__':
     visualizer = Visualizer(opt)
     total_steps = 0
 
-    for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
+    for epoch in range(1, opt.niter + opt.niter_decay + 1):
         epoch_start_time = time.time()
         iter_data_time = time.time()
         epoch_iter = 0
@@ -43,7 +43,7 @@ if __name__ == '__main__':
 
             if total_steps % opt.save_latest_freq == 0:
                 print('saving the latest model (epoch %d, total_steps %d)' %
-                      (epoch, total_steps))
+                      (opt.epoch_count + epoch, total_steps))
                 model.save('latest')
 
             iter_data_time = time.time()
@@ -51,8 +51,8 @@ if __name__ == '__main__':
             print('saving the model at the end of epoch %d, iters %d' %
                   (epoch, total_steps))
             model.save('latest')
-            model.save(epoch)
+            model.save(opt.epoch_count + epoch)
 
         print('End of epoch %d / %d \t Time Taken: %d sec' %
-              (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
+              (opt.epoch_count + epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
         model.update_learning_rate()