OrionZheng
commited on
Commit
•
ab59aaf
1
Parent(s):
b021c89
Update modeling_openmoe.py
Browse files- modeling_openmoe.py +1 -1
modeling_openmoe.py
CHANGED
@@ -371,7 +371,7 @@ class OpenMoeAttention(nn.Module):
|
|
371 |
value_states = repeat_kv(value_states, self.num_key_value_groups)
|
372 |
|
373 |
if HAS_FLASH_ATTN and use_kernel:
|
374 |
-
from flash_attn import flash_attn_func
|
375 |
|
376 |
query_states = query_states.transpose(1, 2)
|
377 |
key_states = key_states.transpose(1, 2)
|
|
|
371 |
value_states = repeat_kv(value_states, self.num_key_value_groups)
|
372 |
|
373 |
if HAS_FLASH_ATTN and use_kernel:
|
374 |
+
exec("from flash_attn import flash_attn_func")
|
375 |
|
376 |
query_states = query_states.transpose(1, 2)
|
377 |
key_states = key_states.transpose(1, 2)
|