diff --git a/structures/Tensorflow.js b/structures/Tensorflow.js
index 218e18ad..50130c5f 100644
--- a/structures/Tensorflow.js
+++ b/structures/Tensorflow.js
@@ -78,8 +78,10 @@ module.exports = class Tensorflow {
 		const styleTensor = tf.node.decodeImage(styleImg, 3);
 		const loadedStyle = styleTensor.div(tf.scalar(255)).expandDims();
 		styleTensor.dispose();
-		const stylePrediction = await this.styleModel.predict(loadedStyle);
+		const reshapedStyle = loadedStyle.reshape([1, 1, 1, 100]);
 		loadedStyle.dispose();
+		const stylePrediction = await this.styleModel.predict(reshapedStyle);
+		reshapedStyle.dispose();
 		const stylizedImage = await this.transformerModel.predict([loadedImage, stylePrediction.squeeze()]);
 		loadedImage.dispose();
 		stylePrediction.dispose();