Fix bug

ROCm · Jul 18, 2024 · bfbc3ef · bfbc3ef
1 parent 4787d45
commit bfbc3ef
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/python/perf-kernels/flash-attention.py b/python/perf-kernels/flash-attention.py
@@ -1171,6 +1171,7 @@ def test_op_fwd_bias(Z, H, N_CTX_Q, N_CTX_K, D_HEAD, causal, use_bias, dtype):
     if causal:
         input_metadata.need_causal()
     if use_bias:
+        bias = torch.randn((1, H, N_CTX_Q, N_CTX_K), dtype=dtype, device="cuda")
         input_metadata.need_bias(bias, Z, H, N_CTX_Q, N_CTX_K)
     else:
         bias = None