1 files changed, 3 insertions, 2 deletions
diff --git a/app/relay/modules/pix2pixhd.js b/app/relay/modules/pix2pixhd.js
index 5f7b1c7..93f9c3a 100644
--- a/app/relay/modules/pix2pixhd.js
+++ b/app/relay/modules/pix2pixhd.js
@@ -64,7 +64,7 @@ const train = {
       '--name', dataset,
       '--model', 'pix2pixHD',
       '--label_nc', 0, '--no_instance',
-      '--niter', task.epochs,
+      '--niter', task.epochs || 1,
       '--niter_decay', 0,
       '--save_epoch_freq', 1,
     ]
@@ -243,7 +243,8 @@ const splice = {
       '--folder_id', task.opt.folder_id,
       '--endpoint', process.env.API_REMOTE + '/api/file/',
     ]
-  }
+  },
+  after: 'train'
 }
 
 function pad(num, size) {