continue train

1 files changed, 2 insertions, 1 deletions

diff --git a/app/relay/modules/pix2pixhd.js b/app/relay/modules/pix2pixhd.js
index a01e321..db40781 100644
--- a/app/relay/modules/pix2pixhd.js
+++ b/app/relay/modules/pix2pixhd.js
@@ -49,8 +49,9 @@ const train = {
     const datasets_path = path.join(cwd, 'datasets', dataset)
     const checkpoints_path = path.join(cwd, 'checkpoints', dataset)
     const iter_txt = path.join(checkpoints_path, 'iter.txt')
+    const checkpoint_path = path.join(checkpoints_path, 'latest_net_G.pth')
     console.log(dataset, iter_txt)
-    if (fs.existsSync(iter_txt)) {
+    if (fs.existsSync(iter_txt) && fs.existsSync(checkpoint_path)) {
       const iter = fs.readFileSync(iter_txt).toString().split('\n');
       console.log(iter)
       epoch = iter[0] || 0