Update train.py

yawudede · Nov 9, 2018 · 9346db4 · 9346db4
1 parent fe3bc52
commit 9346db4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/train.py b/train.py
@@ -42,6 +42,12 @@
                 if opt.display_id > 0:
                     visualizer.plot_current_losses(epoch, float(epoch_iter) / dataset_size, opt, losses)
 
+            if opt.save_by_iter:
+                if total_steps % opt.save_iter_freq == 0:
+                    print('saving model at: (iters %d)' %
+                          (total_steps))
+                    model.save('iter_'+str(total_steps))
+
             if total_steps % opt.save_latest_freq == 0:
                 print('saving the latest model (epoch %d, total_steps %d)' %
                       (epoch, total_steps))