handle checkpoint does not exist

lucidrains · Jan 16, 2021 · 8b1e707 · 8b1e707
1 parent dd235e8
commit 8b1e707
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/lightweight_gan/lightweight_gan.py b/lightweight_gan/lightweight_gan.py
@@ -1353,8 +1353,10 @@ def load(self, num=-1, print_version=True):
         name = num
         if num == -1:
             checkpoints = self.get_checkpoints()
-            name = checkpoints[-1]
-            print(f'continuing from previous epoch - {name}')
+
+            if exists(checkpoints):
+                name = checkpoints[-1]
+                print(f'continuing from previous epoch - {name}')
 
         self.steps = name * self.save_every
 
@@ -1379,6 +1381,6 @@ def get_checkpoints(self):
         saved_nums = sorted(map(lambda x: int(x.stem.split('_')[1]), file_paths))
 
         if len(saved_nums) == 0:
-            return
+            return None
 
         return saved_nums
diff --git a/lightweight_gan/version.py b/lightweight_gan/version.py
@@ -1 +1 @@
-__version__ = '0.17.2'
+__version__ = '0.17.3'