ソースを参照

fix backward for when query and key have different contiguity (#818)

Brian Hirsh 1 年間 前
コミット
2423cca3ad
1 ファイル変更1 行追加1 行削除
  1. 1 1
      csrc/flash_attn/flash_api.cpp

+ 1 - 1
csrc/flash_attn/flash_api.cpp

@@ -830,7 +830,7 @@ mha_bwd(const at::Tensor &dout,  // batch_size x seqlen_q x num_heads, x head_si
         TORCH_CHECK(dv.stride(-1) == 1, "dv must have contiguous last dimension");
         CHECK_SHAPE(dv, batch_size, seqlen_k, num_heads_k, head_size);
     } else {
-        dv = torch::empty_like(k);
+        dv = torch::empty_like(v);
     }
 
     at::Tensor dout_padded;