save

ROCm · Oct 16, 2024 · ba7e5d9 · ba7e5d9
1 parent 6cc25ba
commit ba7e5d9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/flash_attn/flash_attn_triton_amd/bwd_ref.py b/flash_attn/flash_attn_triton_amd/bwd_ref.py
@@ -1,7 +1,7 @@
 import torch
 import math
 
-DEBUG=False
+DEBUG = False
 
 def attention_backward_pytorch_ref_impl(do, q, k, v, o, softmax_lse, sm_scale, causal, layout, use_exp2, bwd_preprocessing_use_o):
     # ensure the layout is 'bhsd'

diff --git a/flash_attn/flash_attn_triton_amd/fwd_ref.py b/flash_attn/flash_attn_triton_amd/fwd_ref.py
@@ -1,7 +1,7 @@
 import math
 import torch
 
-DEBUG=False
+DEBUG = False
 
 def attention_forward_pytorch_ref_impl(q, k, v, sm_scale, causal, layout, use_exp2):
     """compute reference output and softmax_lse using PyTorch's built-in function"""

diff --git a/tests/test_flash_attn_triton.py b/tests/test_flash_attn_triton.py
@@ -19,7 +19,7 @@
 from flash_attn.flash_attn_interface import _get_block_size_n
 from flash_attn.layers.rotary import apply_rotary_emb
 
-DEBUG=False
+DEBUG = False
 # Test ROCM Triton Backend
 USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_USE_TRITON_ROCM", "FALSE") == "TRUE"
 if USE_TRITON_ROCM: