1 files changed, 1 insertions, 1 deletions
diff --git a/main.py b/main.py
index db620a0a..89aa03d9 100644
--- a/main.py
+++ b/main.py
@@ -62,7 +62,7 @@ def main():
     # my_model.vgg16.load_weights(ARGS.load_vgg, by_name=True)
     train(my_model)
     
-    final_image = my_model.x
+    final_image = tf.squeeze(my_model.x)
 
     plt.imshow(final_image)