Update main.py

main.py

@@ -105,8 +105,9 @@ class UnifiedModel(nn.Module):
                 input_ids=input_id,
                 attention_mask=attn_mask
             )
-            hidden_states.append(outputs.logits)
-        concatenated_hidden_states = torch.cat(hidden_states, dim=-1)
+            hidden_states.append(outputs.pooler_output) # Usar pooler_output para obtener un vector de 768
+
+        concatenated_hidden_states = torch.cat(hidden_states, dim=1) # Concatenar en la dimensión correcta
         logits = self.classifier(concatenated_hidden_states)
         return logits
 
@@ -119,7 +120,8 @@ class UnifiedModel(nn.Module):
             model.load_state_dict(torch.load(model_data_bytes))
         else:
             model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
-        return UnifiedModel([model])
+        
+        return UnifiedModel([model, model]) # Asegurar que se usa una lista de modelos, en este caso 2
 
 # Dataset para entrenamiento
 class SyntheticDataset(Dataset):