1 files changed, 2 insertions, 0 deletions
diff --git a/main.py b/main.py
index 063670b8..8363fcef 100644
--- a/main.py
+++ b/main.py
@@ -55,7 +55,9 @@ def main():
     print('this is',ARGS.content)
 
     content_image = imread(ARGS.content)
+    content_image = np.resize(content_image, (255, 255, 3))
     style_image = imread(ARGS.style)
+    style_image = np.resize(style_image, (255, 255, 3))
     my_model = YourModel(content_image=content_image, style_image=style_image)
     train(my_model)