Update models/unet_gray2color.py

models/unet_gray2color.py

@@ -26,9 +26,14 @@ class SelfAttentionLayer(Layer):
         self.ln = LayerNormalization(epsilon=1e-6)
 
     def build(self, input_shape):
-        # Initialize the MultiHeadAttention and LayerNormalization layers
-        self.mha.build(input_shape)  # Build MultiHeadAttention
-        self.ln.build(input_shape)   # Build LayerNormalization
+        # input_shape: (batch_size, height, width, channels)
+        # For self-attention, query, key, and value have the same shape
+        batch_size, height, width, channels = input_shape
+        attention_shape = (batch_size, height * width, channels)  # Shape after reshape
+        # Build MultiHeadAttention with query, key, and value shapes
+        self.mha.build(query_shape=attention_shape, value_shape=attention_shape)
+        # Build LayerNormalization with the original input shape
+        self.ln.build(input_shape)
         super(SelfAttentionLayer, self).build(input_shape)
 
     def call(self, x):