1 files changed, 4 insertions, 1 deletions
diff --git a/main.py b/main.py
index 4497cce0..193d7c7a 100644
--- a/main.py
+++ b/main.py
@@ -70,7 +70,10 @@ def main():
     
     final_image = tf.squeeze(my_model.x)
 
-    plt.imshow(final_image).astype('uint8')
+    # convert image to uint8
+    final_image = tf.cast(final_image, tf.uint8)
+
+    plt.imshow(final_image)
 
     imsave(ARGS.savefile, final_image)