fix bugs

dnaflash.py

@@ -351,8 +351,9 @@ class FLASHTransformer(nn.Module):
 
         for flash in self.layers:
             x = flash(x, mask = mask)
-
-        return self.to_logits(x), x
+        x_norm = self.to_logits[0](x)
+        logits = self.to_logits[1](x_norm)
+        return logits, x_norm
 
 class FLASHTransformerConfig(PretrainedConfig):
     model_type = "flash_transformer"