Merge pull request PaddlePaddle#202 from fedebotu/main

[BugFix] avoid bug on ImportError
AnnaTrainingG · May 6, 2023 · 5bf7f57 · 5bf7f57
2 parents a9a4b4e + 69f5f7d
commit 5bf7f57
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/flash_attn/models/gpt.py b/flash_attn/models/gpt.py
@@ -46,7 +46,7 @@
 try:
     from flash_attn.ops.rms_norm import RMSNorm, dropout_add_rms_norm
 except ImportError:
-    RMSNorm, dropout_add_rms_norm = None
+    RMSNorm, dropout_add_rms_norm = None, None
 
 try:
     from flash_attn.ops.rms_norm import dropout_add_rms_norm_parallel_residual

diff --git a/flash_attn/modules/block.py b/flash_attn/modules/block.py
@@ -26,7 +26,7 @@
 try:
     from flash_attn.ops.rms_norm import RMSNorm, dropout_add_rms_norm
 except ImportError:
-    RMSNorm, dropout_add_rms_norm = None
+    RMSNorm, dropout_add_rms_norm = None, None
 
 try:
     from flash_attn.ops.rms_norm import dropout_add_rms_norm_parallel_residual