Update attention.py

attention.py

@@ -87,7 +87,7 @@ def flash_attn_fn(query, key, value, n_heads, softmax_scale=None, attn_bias=None
 
 def triton_flash_attn_fn(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
     try:
-        from .flash_attn_triton import flash_attn_func
+        from flash_attn import flash_attn_triton
     except:
         raise RuntimeError('Please install flash-attn==1.0.3.post0 and triton==2.0.0.dev20221202')
     check_valid_inputs(query, key, value)