Update Model.py

Model.py

@@ -153,7 +153,6 @@ class LoU(layers.Layer):
         self.Vr = Lo(d_model)
         self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
         self.O = layers.Dense(d_model, dtype='float32')
-        self.glu = SwiGLU(d_model)
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         self.norm1 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         
@@ -201,16 +200,26 @@ class LoU(layers.Layer):
         score_clipped = tf.clip_by_value(score_norm, -self.clip_value, self.clip_value)
         x_comb = score_clipped * V
         out = self.proj(x_comb)
-        out = self.glu(out)
         out = self.norm(out + residual)
         return tf.cast(out, x.dtype)
 
+class Block(layers.Layer):
+    def __init__(self, d_model):
+        super().__init__()
+        self.lou = LoU(d_model)
+        self.glu = SwiGLU(d_model)
+
+    def call(self, x):
+        x = self.lou(x)
+        x = self.glu(x)
+        return x
+
 class ReLM(tf.keras.Model):
     def __init__(self, vocab_size, max_seq_len, d_model, n_layers, dropout_rate=0.1):
         super().__init__()
         self.token_embedding = layers.Embedding(vocab_size, d_model)
         self.pos_embedding = layers.Embedding(max_seq_len, d_model)
-        self.blocks = [LoU(d_model) for _ in range(n_layers)]
+        self.blocks = [Block(d_model) for _ in range(n_layers)]
         self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype="float32")
 
     def call(self, x, training=False):