From 73d69a46d421fd68f3153d016d370b1a142ba1c8 Mon Sep 17 00:00:00 2001
From: Dragon Fire <danielbodendahl@gmail.com>
Date: Wed, 8 May 2024 18:08:42 -0400
Subject: [PATCH] Fix

---
 structures/Tensorflow.js | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

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();