Update modeling_motif.py

modeling_motif.py

@@ -811,9 +811,8 @@ class MotifDecoderLayer(nn.Module):
         self.self_attn = MOTIF_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
         self.mlp = MotifMLP(config)
                 
-        RMSNorm = MorehRMSNorm
-        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.input_layernorm = MotifRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = MotifRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
 
     def forward(
@@ -1048,8 +1047,7 @@ class MotifModel(MotifPreTrainedModel):
         num_hidden_layers = config.num_hidden_layers
         self.layers = nn.ModuleList([MotifDecoderLayer(config = config, layer_idx=layer_idx) for layer_idx in range(num_hidden_layers)])
         self._attn_implementation = config._attn_implementation
-        RMSNorm = MorehRMSNorm
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm = MotifRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hidden_size = config.hidden_size
         self.num_heads = config.num_attention_heads
         self.head_dim = self.hidden_size // self.num_heads