Fix docs shape

Author: EricLBuehler

candle-nn/src/cpu_flash_attention.rs

@@ -71,7 +71,7 @@ fn vec_dot<T: WithDType + Sum + Copy + std::ops::Mul<Output = T>>(a: &[T], b: &[
 ///
 /// - This supports ALiBi with `max_bias` as well as softcapping with `softcap`.
 ///
-/// **Output shape:** (bs, seq, qhead, v_hidden)
+/// **Output shape:** (bs, qhead, seq, v_hidden)
 pub fn run_flash_attn_cpu<T>(
     q: &Tensor,
     k: &Tensor,

candle-nn/tests/cpu_flash_attn.rs

@@ -30,8 +30,7 @@ fn cpu_flash_attn() -> Result<()> {
         softmax_scale,
         None,
         None,
-    )?
-    .transpose(1, 2)?;
+    )?;
 
     let out_arr: Vec<f32> = out.flatten_all()?.to_vec1()?;
     let ground_truth_arr: Vec<f32> = ground_truth.flatten_all()?.to_vec1()?;