augment plus 2

1 files changed, 2 insertions, 2 deletions

diff --git a/augment.py b/augment.py
index b01f04c..4aab5f1 100755
--- a/augment.py
+++ b/augment.py
@@ -67,7 +67,7 @@ if data_opt.tag == '':
 else:
   tag = data_opt.tag
 
-if opt.current_epoch == 'latest':
+if opt.which_epoch == 'latest':
   iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
   if os.path.exists(iter_path):
     try:
@@ -78,7 +78,7 @@ if opt.current_epoch == 'latest':
   else:    
     current_epoch, epoch_iter = 1, 0
 else:
-  current_epoch = opt.current_epoch
+  current_epoch = opt.which_epoch
 
 epoch_id = "{:02d}_{:04d}_{:04d}".format(current_epoch, data_opt.augment_take, data_opt.augment_make)
 opt.render_dir = os.path.join(opt.results_dir, opt.name, epoch_id)