Update Mo.py

Mo.py

@@ -112,18 +112,6 @@ dataset = dataset.shuffle(2000, seed=SEED).batch(batch_size, drop_remainder=True
 with strategy.scope():
     dist_dataset = strategy.experimental_distribute_dataset(dataset)
 
-class SwiGLU(layers.Layer):
-    def __init__(self, d_model):
-        super().__init__()
-        self.W = layers.Dense(3500, dtype='float32')
-        self.W1 = layers.Dense(d_model, dtype='float32')
-    def call(self, x):
-        x = tf.cast(x, tf.float32)
-        x = self.W(x)
-        a, b = tf.split(x, 2, axis=-1)
-        out = self.W1(tf.nn.silu(a) * b)     
-        return tf.cast(out, x.dtype)
-
 class SwiGLU(layers.Layer):
     def __init__(self, d_model, d_ff):
         super().__init__()
@@ -200,13 +188,10 @@ class Block(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
         self.lou = LoU(d_model)
-        self.glu = SwiGLU(d_model)
-        self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         self.lo = Lo(d_model)
 
     def call(self, x):
         x = self.lou(x)
-        x = self.norm(self.glu(x)) + x
         x = self.lo(x)
         return x