changes

1 files changed, 3 insertions, 0 deletions

diff --git a/main.py b/main.py
index 85ee4862..c36658f3 100644
--- a/main.py
+++ b/main.py
@@ -72,6 +72,9 @@ def main():
     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)
+
     plt.imshow(final_image)
 
     imsave(ARGS.savefile, final_image)