Update main.py

main.py

@@ -96,7 +96,8 @@ class UnifiedModel(nn.Module):
         super(UnifiedModel, self).__init__()
         self.models = nn.ModuleList(models)
         hidden_size = self.models[0].config.hidden_size
-        self.classifier = nn.Linear(len(models) * hidden_size, 3) # 3 clases
+        self.projection = nn.Linear(len(models) * 3, hidden_size)  # Nueva capa lineal para proyección
+        self.classifier = nn.Linear(hidden_size, 3) # 3 clases
 
     def forward(self, input_ids, attention_mask):
         hidden_states = []
@@ -108,7 +109,8 @@ class UnifiedModel(nn.Module):
             hidden_states.append(outputs.logits)  # Usar directamente outputs.logits
 
         concatenated_hidden_states = torch.cat(hidden_states, dim=1)
-        logits = self.classifier(concatenated_hidden_states)
+        projected_features = self.projection(concatenated_hidden_states) # Proyectar a hidden_size
+        logits = self.classifier(projected_features) 
         return logits
 
     @staticmethod