p2p task

1 files changed, 7 insertions, 3 deletions

diff --git a/app/relay/modules/pix2pix.js b/app/relay/modules/pix2pix.js
index 6979be5..0e67621 100644
--- a/app/relay/modules/pix2pix.js
+++ b/app/relay/modules/pix2pix.js
@@ -71,7 +71,7 @@ const train = {
         console.log(task.module, task.dataset, epoch, task.epochs)
       } catch (e) { }
     }
-    return [
+    let args = [
       '--dataroot', datasets_path,
       '--module-name', task.module,
       '--name', task.dataset,
@@ -85,12 +85,16 @@ const train = {
       '--epoch_count', task.epoch + task.epochs + 1,
       '--niter', task.epochs,
       '--niter_decay', 0,
-      '--which_epoch', 'latest',
-      '--continue_train',
       '--no_lsgan',
       '--norm', 'batch',
       '--pool_size', '0',
     ]
+    if (epoch) {
+      args = args.concat([
+        '--which_epoch', 'latest',
+        '--continue_train',
+      })
+    }
   },
 }
 const generate = {