fix: flash attention import

jina-ai · Nov 10, 2022 · 5cedb3d · 5cedb3d
1 parent c9383cf
commit 5cedb3d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/server/clip_server/model/model.py b/server/clip_server/model/model.py
@@ -27,9 +27,10 @@
 from open_clip.factory import _MODEL_CONFIGS
 
 # Use flash attention
-FLASH_ATTENTION_AVAILABLE = True
 try:
     from clip_server.model.flash_attention import MultiheadAttention
+
+    FLASH_ATTENTION_AVAILABLE = True
 except:
     FLASH_ATTENTION_AVAILABLE = False