OrionZheng commited on
Commit
ab59aaf
1 Parent(s): b021c89

Update modeling_openmoe.py

Browse files
Files changed (1) hide show
  1. 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)