Update dnaflash.py

dnaflash.py

@@ -408,7 +408,7 @@ class FLASHTransformerForPretrained(PreTrainedModel):
             reduce_group_non_causal_attn=config.reduce_group_non_causal_attn
         )
 
-    def forward(self, input_ids, mask=None):
-        logits, x = self.model(input_ids, mask=mask)
+    def forward(self,inputs):
+        logits, x = self.model(inputs["input_ids"], mask=inputs["attention_mask"])
         return MaskedLMOutput(logits=logits, hidden_states=x, loss=None, attentions=None)