Update model.py

model.py

@@ -53,7 +53,6 @@ class BERTMultiAttentionModel(PreTrainedModel):
         super(BERTMultiAttentionModel, self).__init__(config)
         self.config = config
 
-        # Initialize the transformer model
         self.transformer = AutoModel.from_pretrained(config.transformer)
         self.cross_attention = MultiHeadAttention(config)
         self.fc1 = nn.Linear(config.hidden_size * 2, 256)
@@ -77,7 +76,7 @@ class BERTMultiAttentionModel(PreTrainedModel):
         combined_output = self.dropout1(torch.relu(combined_output))
         combined_output = combined_output.unsqueeze(1)
         _, (hidden_state, _) = self.rnn(combined_output)
-        hidden_state_concat = torch.cat([hidden_state[0], hidden_state[1]], dim=-1)
+        hidden_state_concat = torch.cat([hidden_state[-2], hidden_state[-1]], dim=-1)
         hidden_state_proj = self.layer_norm_proj(self.fc_proj(hidden_state_concat))
         hidden_state_proj = self.dropout2(hidden_state_proj)
         final = self.fc_final(hidden_state_proj)