Merge branch 'master' of https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix

1 files changed, 4 insertions, 4 deletions

diff --git a/train.py b/train.py
index da497ce..ae129b9 100644
--- a/train.py
+++ b/train.py
@@ -8,8 +8,8 @@ from util.visualizer import Visualizer
 
 data_loader = CreateDataLoader(opt)
 dataset = data_loader.load_data()
-num_train = len(data_loader)
-print('#training images = %d' % num_train)
+dataset_size = len(data_loader)
+print('#training images = %d' % dataset_size)
 
 model = create_model(opt)
 visualizer = Visualizer(opt)
@@ -21,7 +21,7 @@ for epoch in range(1, opt.niter + opt.niter_decay + 1):
     for i, data in enumerate(dataset):
         iter_start_time = time.time()
         total_steps += opt.batchSize
-        epoch_iter = total_steps % num_train
+        epoch_iter = total_steps - dataset_size * (epoch - 1)
         model.set_input(data)
         model.optimize_parameters()
 
@@ -32,7 +32,7 @@ for epoch in range(1, opt.niter + opt.niter_decay + 1):
             errors = model.get_current_errors()
             visualizer.print_current_errors(epoch, epoch_iter, errors, iter_start_time)
             if opt.display_id > 0:
-                visualizer.plot_current_errors(epoch, float(epoch_iter)/num_train, opt, errors)
+                visualizer.plot_current_errors(epoch, float(epoch_iter)/dataset_size, opt, errors)
 
         if total_steps % opt.save_latest_freq == 0:
             print('saving the latest model (epoch %d, total_steps %d)' %