Update dnaflash.py

dnaflash.py

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