Update Model.py

Model.py

@@ -158,7 +158,6 @@ class LoSoU(layers.Layer):
         self.K = layers.Dense(96, dtype='float32')
         self.V = Lo(d_model)  # Lo already handles casting to model dtype; we'll cast back to float32
         self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
-        self.O = layers.Dense(d_model, dtype='float32')
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
 
         # 동적 alpha 계산을 위한 레이어
@@ -242,17 +241,7 @@ class LoSoU(layers.Layer):
         x_comb = score_clipped * V  # (B, L, d_model)
 
         out = self.proj(x_comb)  # (B, L, d_model)
-
-        # ensure out dim even for split
-        d = out.shape[-1]  # this is an int (static shape)
-        if d is not None and d % 2 == 1:
-            out = tf.pad(out, [[0,0],[0,0],[0,1]])
-
-        a, b = tf.split(out, 2, axis=-1)
-        gated = tf.nn.silu(a) * b
-        out = self.O(gated)
-
-        out = self.norm(out + residual)
+        out = self.norm(out)
 
         # cast back to original dtype for downstream layers
         return tf.cast(out, x.dtype)
@@ -271,10 +260,9 @@ class ReLaM(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, 128)
-        self.pos_embedding = layers.Embedding(max_seq_len, d_model)
+        self.pos_embedding = layers.Embedding(max_seq_len, 128)
         self.blocks = [Block(d_model, hyper_n=1) for _ in range(n_layers)]
         self.proj = layers.Dense(128)
-        # LayerNormalization은 float32로 해서 정밀도 문제 방지
         self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype="float32")
 
     def call(self, x, training=False):