FIX RETRAIN LOGIC

1 files changed, 7 insertions, 7 deletions

diff --git a/app/relay/modules/pix2pixhd.js b/app/relay/modules/pix2pixhd.js
index 02431e5..3613e1c 100644
--- a/app/relay/modules/pix2pixhd.js
+++ b/app/relay/modules/pix2pixhd.js
@@ -45,22 +45,22 @@ const train = {
   script: 'train.py',
   params: (task) => {
     let epoch = 0
-
-    const datasets_path = path.join(cwd, 'datasets', task.dataset)
-    const checkpoints_path = path.join(cwd, 'checkpoints', task.dataset)
+    const dataset = task.dataset.toLowerCase()
+    const datasets_path = path.join(cwd, 'datasets', dataset)
+    const checkpoints_path = path.join(cwd, 'checkpoints', dataset)
     if (fs.existsSync(checkpoints_path)) {
       try {
         const checkpoints = fs.readdirSync(checkpoints_path)
-        checkpoints.forEach(c => {
-          epoch = Math.max(parseInt(c.name) || 0, epoch)
+        checkpoints.forEach(name => {
+          epoch = Math.max(parseInt(name) || 0, epoch)
         })
-        console.log(task.module, task.dataset, epoch, task.epochs)
+        console.log(task.module, dataset, epoch, task.epochs)
       } catch (e) { }
     }
     let args = [
       '--dataroot', datasets_path,
       '--module_name', task.module,
-      '--name', task.dataset,
+      '--name', dataset,
       '--model', 'pix2pixHD',
       '--label_nc', 0, '--no_instance',
       '--niter', task.epochs,