Update modeling_neobert.py

modeling_neobert.py

@@ -206,7 +206,7 @@ class EncoderBlock(nn.Module):
         self.ffn = SwiGLU(config.hidden_size, intermediate_size, config.hidden_size, bias=False)
 
         # Layer norms
-        rms_norm_cls = nn.RMSNorm if config._attn_implementation != 'onnx_eager' else NeoBERTEagerRMSNorm
+        rms_norm_cls = nn.RMSNorm if config._attn_implementation != 'onnx_eager' and hasattr(nn, 'RMSNorm') else NeoBERTEagerRMSNorm
         self.attention_norm = rms_norm_cls(config.hidden_size, config.norm_eps)
         self.ffn_norm = rms_norm_cls(config.hidden_size, config.norm_eps)
 
@@ -315,7 +315,7 @@ class NeoBERT(NeoBERTPreTrainedModel):
         for _ in range(config.num_hidden_layers):
             self.transformer_encoder.append(EncoderBlock(config))
 
-        rms_norm_cls = nn.RMSNorm if config._attn_implementation != 'onnx_eager' else NeoBERTEagerRMSNorm    
+        rms_norm_cls = nn.RMSNorm if config._attn_implementation != 'onnx_eager' and hasattr(nn, 'RMSNorm') else NeoBERTEagerRMSNorm
         self.layer_norm = rms_norm_cls(config.hidden_size, config.norm_eps)
 
         # Initialize weights and apply final processing