fix issues with erf and xavier init

modeling_siglip.py

@@ -95,7 +95,12 @@ def _trunc_normal_(tensor, mean, std, a, b):
 
     # Use inverse cdf transform for normal distribution to get truncated
     # standard normal
-    tensor.erfinv_()
+    if tensor.dtype == torch.bfloat16:
+        tensor = tensor.to(torch.float32)
+        tensor.erfinv_()
+        tensor = tensor.to(torch.bfloat16)
+    else:
+        tensor.erfinv_()
 
     # Transform to proper mean, std
     tensor.mul_(std * math.sqrt(2.0))
@@ -670,6 +675,7 @@ class SiglipPreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         """Initialize the weights"""
+        
         if isinstance(module, SiglipVisionEmbeddings):
             width = (
                 self.config.vision_config.hidden_size
@@ -680,22 +686,22 @@ class SiglipPreTrainedModel(PreTrainedModel):
         elif isinstance(module, nn.Embedding):
             default_flax_embed_init(module.weight)
         elif isinstance(module, SiglipAttention):
-            nn.init.xavier_uniform_(module.q_proj.weight)
-            nn.init.xavier_uniform_(module.k_proj.weight)
-            nn.init.xavier_uniform_(module.v_proj.weight)
-            nn.init.xavier_uniform_(module.out_proj.weight)
+            nn.init.normal_(module.q_proj.weight)
+            nn.init.normal_(module.k_proj.weight)
+            nn.init.normal_(module.v_proj.weight)
+            nn.init.normal_(module.out_proj.weight)
             nn.init.zeros_(module.q_proj.bias)
             nn.init.zeros_(module.k_proj.bias)
             nn.init.zeros_(module.v_proj.bias)
             nn.init.zeros_(module.out_proj.bias)
         elif isinstance(module, SiglipMLP):
-            nn.init.xavier_uniform_(module.fc1.weight)
-            nn.init.xavier_uniform_(module.fc2.weight)
+            nn.init.normal_(module.fc1.weight)
+            nn.init.normal_(module.fc2.weight)
             nn.init.normal_(module.fc1.bias, std=1e-6)
             nn.init.normal_(module.fc2.bias, std=1e-6)
         elif isinstance(module, SiglipMultiheadAttentionPoolingHead):
-            nn.init.xavier_uniform_(module.probe.data)
-            nn.init.xavier_uniform_(module.attention.in_proj_weight.data)
+            nn.init.normal_(module.probe.data)
+            nn.init.normal_(module.attention.in_proj_weight.data)
             nn.init.zeros_(module.attention.in_proj_bias.data)
         elif isinstance(module, SiglipModel):
             logit_scale_init = torch.log(torch.tensor(1.0))