Update modeling_intern_vit.py

modeling_intern_vit.py

@@ -129,6 +129,12 @@ except Exception:
     pass
 
 
+NORM2FN = {
+    'rms_norm': InternRMSNorm,
+    'layer_norm': nn.LayerNorm,
+}
+
+
 class InternVisionEmbeddings(nn.Module):
     def __init__(self, config: InternVisionConfig):
         super().__init__()
@@ -267,11 +273,12 @@ class InternVisionEncoderLayer(nn.Module):
         super().__init__()
         self.embed_dim = config.hidden_size
         self.intermediate_size = config.intermediate_size
+        self.norm_type = config.norm_type
 
         self.attn = InternAttention(config)
         self.mlp = InternMLP(config)
-        self.norm1 = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
-        self.norm2 = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
+        self.norm2 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
 
         self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
         self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))