Update dnaflash.py

dnaflash.py

@@ -475,7 +475,7 @@ class FLASHTransformerForSequenceClassification(FLASHTransformerForPretrained):
         )
         hidden_states = outputs["hidden_states"]
         input_mask_expanded = input_ids["attention_mask"].unsqueeze(-1).expand(hidden_states.size())  # 维度匹配
-        mean_pooled = torch.sum(token_embeddings * input_mask_expanded, dim=1) / input_mask_expanded.sum(dim=1)  # 计算加权平均
+        mean_pooled = torch.sum(hidden_states * input_mask_expanded, dim=1) / input_mask_expanded.sum(dim=1)  # 计算加权平均
         logits = self.score(mean_pooled)
 
         loss = None