Update modeling_custom.py

modeling_custom.py

@@ -38,14 +38,13 @@ class GatingNetwork(nn.Module):
         dropout_rate = dropout
         for i in range(n_hidden):
             layers.append(nn.Linear(in_features, hidden_dim, bias=False)) # for BN
-            nn.init.kaiming_normal_(layers[-1].weight, mode='fan_in', nonlinearity='relu')
+            #nn.init.kaiming_normal_(layers[-1].weight, mode='fan_in', nonlinearity='relu')
             layers.append(nn.ReLU())
             layers.append(nn.BatchNorm1d(hidden_dim))
             if dropout_rate > 0 and i < n_hidden - 1: # no dropout before last layer for more stability and precision
                 layers.append(nn.Dropout(dropout_rate))
 
             in_features = hidden_dim
-        # self.dropout_list = [nn.Dropout(dropout_rate), nn.Dropout(dropout_rate)]
         layers.append(nn.Linear(in_features, out_features, bias=bias))
         self.layers = nn.ModuleList(layers)
         # print("Gating network layers:", self.layers)
@@ -117,7 +116,7 @@ class Gemma2ForQuantileSequenceClassification(Gemma2PreTrainedModel):
 
         # Initialize weights and apply final processing
         self.gating = GatingNetwork(config.hidden_size, self.num_objectives,
-                                    temperature=config_dict.get("gating_temperature", 1),
+                                    temperature=config_dict.get("gating_temperature", 2),
                                     hidden_dim=config_dict.get("gating_hidden_dim", 1024),
                                     n_hidden=config_dict.get("gating_n_hidden", 3))