adding vgg weights

1 files changed, 2 insertions, 1 deletions

diff --git a/main.py b/main.py
index 0d707b81..de769baa 100644
--- a/main.py
+++ b/main.py
@@ -53,7 +53,8 @@ def main():
     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)
-    # my_model.vgg16.load_weights(ARGS.load_vgg, by_name=True)
+    my_model.vgg16.build([1, 255, 255, 3])
+    my_model.vgg16.load_weights('vgg16_imagenet.h5', by_name=True)
     train(my_model)
     
     final_image = tf.squeeze(my_model.x)