saving checkpoints

1 files changed, 10 insertions, 0 deletions

diff --git a/main.py b/main.py
index 45714611..87f36013 100644
--- a/main.py
+++ b/main.py
@@ -39,9 +39,17 @@ def parse_args():
 
 def train(model):
     for i in range(hp.num_epochs):
+        if i % 100 == 0:
+            fn = f"checkpoint-images/img-epoch{i}"
+            save_image(fn, model.x)
+            
         print('batch', i)
         model.train_step()
 
+def save_image(filename, image):
+    imsave(filename, image)
+
+
 def main():
     """ Main function. """
     if os.path.exists(ARGS.content):
@@ -66,6 +74,8 @@ def main():
 
     imsave(ARGS.savefile, final_image)
 
+    
+
 
 ARGS = parse_args()
 main()