Update AlphaS2S.py

AlphaS2S.py

@@ -166,13 +166,15 @@ class EncoderBlock(layers.Layer):
         super().__init__()
         self.mha = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)
         self.ffn = SwiGLU(d_model, dff)
+        self.proj = layers.Dense(d_model)
         self.norm1 = layers.LayerNormalization(epsilon=1e-6)
         self.norm2 = layers.LayerNormalization(epsilon=1e-6)
         self.dropout1 = layers.Dropout(dropout)
         self.dropout2 = layers.Dropout(dropout)
     def call(self, x, mask=None, training=False):
+        x = self.proj(x)
         attn_out = self.dropout1(self.mha(x, x, x, attention_mask=mask), training=training)
-        out1 = self.norm1(attn_out)
+        out1 = self.norm1(attn_out + x)
         ffn_out = self.dropout2(self.ffn(out1), training=training)
         return self.norm2(out1 + ffn_out)
 
@@ -182,6 +184,7 @@ class DecoderBlock(layers.Layer):
         self.self_mha = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)
         self.cross_mha = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)
         self.ffn = SwiGLU(d_model, dff)
+        self.proj = layers.Dense(d_model)
         self.norm1 = layers.LayerNormalization(epsilon=1e-6)
         self.norm2 = layers.LayerNormalization(epsilon=1e-6)
         self.norm3 = layers.LayerNormalization(epsilon=1e-6)
@@ -189,8 +192,9 @@ class DecoderBlock(layers.Layer):
         self.dropout2 = layers.Dropout(dropout)
         self.dropout3 = layers.Dropout(dropout)
     def call(self, x, enc_out, training=False):
+        x = self.proj(x)
         attn1 = self.dropout1(self.self_mha(x, x, x, use_causal_mask=True), training=training)
-        out1 = self.norm1(attn1)
+        out1 = self.norm1(attn1 + x)
         attn2 = self.dropout2(self.cross_mha(out1, enc_out, enc_out), training=training)
         out2 = self.norm2(out1 + attn2)
         ffn_out = self.dropout3(self.ffn(out2), training=training)