Update modeling_FalconTST.py

modeling_FalconTST.py

@@ -247,7 +247,7 @@ class RMSNorm(nn.Module):
         return self.weight * hidden_states.to(input_dtype)
 
 
-class FlashAttention(nn.Module):
+class TEDotProductAttention(nn.Module):
     """Implement the scaled dot product attention with softmax.
     Arguments
     ---------
@@ -294,16 +294,6 @@ class FlashAttention(nn.Module):
 
 
 
-class TEDotProductAttention(nn.Module):
-    def __init__(self, flash_attention,):
-        super().__init__()
-        self.flash_attention = flash_attention
-
-    def forward(self, q, k, v, mask=None):
-        # Prioritize using FlashAttention
-        return self.flash_attention(q, k, v, mask)
-
-
 class SelfAttention(nn.Module):
     def __init__(self,config,):
         super().__init__()
@@ -311,9 +301,7 @@ class SelfAttention(nn.Module):
         q_layernorm=config.q_layernorm
         k_layernorm=config.k_layernorm
         self.hidden_size = config.hidden_size
-        self.core_attention = TEDotProductAttention(
-                        flash_attention=FlashAttention(),
-        )
+        self.core_attention = TEDotProductAttention()
         self.linear_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.add_bias_linear,)
         self.linear_qkv =  nn.Linear(self.hidden_size, 3*self.hidden_size, bias=config.add_bias_linear,)
         if q_layernorm: