fixed seed issue

model.py

@@ -90,9 +90,10 @@ class MimicTransformer(Module):
             cls_results = self.model(input_ids, attention_mask=attention_mask, labels=drg_labels, output_attentions=True)
         else:
             cls_results = self.model(input_ids, attention_mask=attention_mask, output_attentions=True)
-        # last_attn = cls_results[-1][-1] # (batch, attn_heads, tokens, tokens)
-        last_attn = torch.mean(torch.stack(cls_results[-1])[:], dim=0)
-        last_layer_attn = torch.mean(last_attn[:, :-3, :, :], dim=1)
+        last_attn = cls_results[-1][-1] # (batch, attn_heads, tokens, tokens)
+        # last_attn = torch.mean(torch.stack(cls_results[-1])[:], dim=0)
+        # last_layer_attn = torch.mean(last_attn[:, :-3, :, :], dim=1)
+        last_layer_attn = last_attn[:, -1, :, :]
         xai_logits = self.linear(last_layer_attn).squeeze(dim=-1)
         return (cls_results, xai_logits)