1 files changed, 3 insertions, 1 deletions
diff --git a/train.py b/train.py
index 7d2a5e9..6dbd66b 100644
--- a/train.py
+++ b/train.py
@@ -20,13 +20,15 @@ for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
 
     for i, data in enumerate(dataset):
         iter_start_time = time.time()
+        visualizer.reset()
         total_steps += opt.batchSize
         epoch_iter += opt.batchSize
         model.set_input(data)
         model.optimize_parameters()
 
         if total_steps % opt.display_freq == 0:
-            visualizer.display_current_results(model.get_current_visuals(), epoch)
+            save_result = total_steps % opt.update_html_freq == 0
+            visualizer.display_current_results(model.get_current_visuals(), epoch, save_result)
 
         if total_steps % opt.print_freq == 0:
             errors = model.get_current_errors()