Update app.py

app.py

@@ -12,11 +12,13 @@ CONTENT_LAYER = [...] # e.g. [('block5_conv4', 1)]
 
 
 def gram_matrix(A):
-    """Compute Gram matrix for style representation."""
-    A = tf.transpose(A, (0, 3, 1, 2))   # (batch, channels, height, width)
-    features = tf.reshape(A, (A.shape[0], A.shape[1], -1))
-    gram = tf.matmul(features, features, transpose_b=True)
-    return gram / tf.cast(tf.shape(features)[-1], tf.float32)
+    """Compute Gram matrix for style representation.
+    A has shape (batch, height, width, channels).
+    """
+    # Flatten height & width into one dimension
+    A = tf.reshape(A, (A.shape[0], -1, A.shape[3]))  # (batch, H*W, C)
+    gram = tf.matmul(A, A, transpose_a=True)         # (batch, C, C)
+    return gram / tf.cast(tf.shape(A)[1], tf.float32)
 
 def compute_content_cost(a_C, a_G):
     """Content cost between content and generated image features."""