1 files changed, 3 insertions, 1 deletions
diff --git a/main.py b/main.py
index 4497cce0..85ee4862 100644
--- a/main.py
+++ b/main.py
@@ -48,6 +48,7 @@ def train(model):
         model.train_step()
 
 def save_image(filename, image):
+    image = transform.resize(image, tf.shape(image), anti_aliasing=True).astype('uint8')
     imsave(filename, image)
 
 
@@ -69,8 +70,9 @@ def main():
     train(my_model)
     
     final_image = tf.squeeze(my_model.x)
+    final_image = transform.resize(final_image, tf.shape(final_image), anti_aliasing=True).astype('uint8')
 
-    plt.imshow(final_image).astype('uint8')
+    plt.imshow(final_image)
 
     imsave(ARGS.savefile, final_image)