Update main.py

main.py

@@ -31,7 +31,8 @@ class UnifiedModel(nn.Module):
     def __init__(self, models):
         super(UnifiedModel, self).__init__()
         self.models = nn.ModuleList(models)
-        self.classifier = nn.Linear(sum([model.config.hidden_size for model in models]), 2)
+        hidden_size = self.models[0].config.hidden_size
+        self.classifier = nn.Linear(len(models) * hidden_size, 2)
 
     def forward(self, input_ids, attention_mask):
         hidden_states = []
@@ -40,7 +41,7 @@ class UnifiedModel(nn.Module):
                 input_ids=input_id,
                 attention_mask=attn_mask
             )
-            hidden_states.append(outputs.logits[:, 0])
+            hidden_states.append(outputs.logits[:, -1, :])
         concatenated_hidden_states = torch.cat(hidden_states, dim=-1)
         logits = self.classifier(concatenated_hidden_states)
         return logits