Update Inference.py

Inference.py

@@ -68,7 +68,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 계산을 위한 레이어
@@ -118,7 +117,7 @@ class LoSoU(layers.Layer):
         # x: (B, L, d_model) maybe bfloat16 or float32
         # cast to float32 for all internal computations
         x_f32 = tf.cast(x, tf.float32)
-        
+        residual = x_f32
 
         # Q, K, V
         q = self.Q(x_f32)   # (B, L, 96)
@@ -133,7 +132,7 @@ class LoSoU(layers.Layer):
         score = g_q * g_k
 
         # 동적 alpha 계산: (B, L, d_model) -> (B, L, 1)
-        alpha_dynamic = self.alpha_linear(x_f32) # (B, L, 1)
+        alpha_dynamic = self.alpha_linear(x_f32) * 0.8 + 0.1 # (B, L, 1)
         # 필요시 alpha_dynamic에 대한 후처리 (예: min/max 등) 가능
         # ex: alpha_dynamic = tf.clip_by_value(alpha_dynamic, 0.01, 0.99)
 
@@ -152,16 +151,6 @@ 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)
 
         # cast back to original dtype for downstream layers