testing

1 files changed, 5 insertions, 2 deletions

diff --git a/main.py b/main.py
index b7d9a1e0..2d8c216a 100644
--- a/main.py
+++ b/main.py
@@ -1,7 +1,9 @@
 import os
 import sys
 import argparse
+
 import tensorflow as tf
+from skimage import transform
 
 import hyperparameters as hp
 from losses import YourModel
@@ -37,7 +39,8 @@ def parse_args():
     return parser.parse_args()
 
 def train(model):
-    for _ in range(hp.num_epochs):
+    for i in range(hp.num_epochs):
+        print('batch', i)
         model.train_step()
 
 def main():
@@ -51,7 +54,7 @@ def main():
 
     content_image = imread(ARGS.content)
     style_image = imread(ARGS.style)
-    style_image = transform.resize(style_image, content_image.shape)
+    style_image = transform.resize(style_image, content_image.shape).astype('uint8')
     my_model = YourModel(content_image=content_image, style_image=style_image)
     my_model.vgg16.build([1, 255, 255, 3])
     my_model.vgg16.load_weights('vgg16_imagenet.h5', by_name=True)