1 files changed, 2 insertions, 0 deletions
diff --git a/main.py b/main.py
index 193d7c7a..c36658f3 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,6 +70,7 @@ 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')
 
     # convert image to uint8
     final_image = tf.cast(final_image, tf.uint8)