Update modeling_afmoe_scm_liger.py

modeling_afmoe_scm_liger.py

@@ -25,7 +25,6 @@ from transformers.integrations import use_kernel_forward_from_hub
 import scattermoe
 
 from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
-from liger_kernel.transformers.rms_norm import LigerRMSNorm as AfmoeSCMRMSNorm
 
 try:
     from .configuration_afmoe_scm import AfmoeSCMConfig
@@ -141,8 +140,8 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     )
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
-#@use_kernel_forward_from_hub("RMSNorm")
-class _AfmoeSCMRMSNorm(nn.Module):
+@use_kernel_forward_from_hub("RMSNorm")
+class AfmoeSCMRMSNorm(nn.Module):
     def __init__(self, hidden_size: int, eps: float):
         """
         AfmoeSCMRMSNorm is equivalent to T5LayerNorm