isdigit

1 files changed, 5 insertions, 1 deletions

diff --git a/augment.py b/augment.py
index 4aab5f1..87d5717 100755
--- a/augment.py
+++ b/augment.py
@@ -80,7 +80,11 @@ if opt.which_epoch == 'latest':
 else:
   current_epoch = opt.which_epoch
 
-epoch_id = "{:02d}_{:04d}_{:04d}".format(current_epoch, data_opt.augment_take, data_opt.augment_make)
+if str(current_epoch).isdigit():
+  epoch_id = "{:02d}_{:04d}_{:04d}".format(int(current_epoch), data_opt.augment_take, data_opt.augment_make)
+else:
+  epoch_id = "{}_{: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)
 
 if not opt.engine and not opt.onnx: