1 files changed, 1 insertions, 1 deletions
diff --git a/main.py b/main.py
index 87f36013..867c3f8d 100644
--- a/main.py
+++ b/main.py
@@ -41,7 +41,7 @@ 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)
+            save_image(fn, tf.squeeze(model.x))
             
         print('batch', i)
         model.train_step()