Fix dk/dv autograd error on TPU flash attention (#8685)

pytorch · Feb 7, 2025 · 0cd1fc2 · 0cd1fc2
1 parent afaf0d0
commit 0cd1fc2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_xla/experimental/custom_kernel.py b/torch_xla/experimental/custom_kernel.py
@@ -449,7 +449,7 @@ def fa_custom_backward(
     if require_grad_ab:
       grad_ab = grads[1]
 
-  if require_grad_k or require_grad_k:
+  if require_grad_k or require_grad_v:
     payload, _ = trace_pallas(
         _flash_attention_bwd_dkv,
         q,