Update main.py

main.py

@@ -105,7 +105,7 @@ class UnifiedModel(nn.Module):
                 input_ids=input_id,
                 attention_mask=attn_mask
             )
-            hidden_states.append(outputs.logits[:, -1, :])  # Obtener el último hidden state
+            hidden_states.append(outputs.logits)  # Usar directamente outputs.logits
 
         concatenated_hidden_states = torch.cat(hidden_states, dim=1)
         logits = self.classifier(concatenated_hidden_states)