Update thunder/tests/test_cudnn_executor.py

Co-authored-by: Ivan Yashchuk <[email protected]>
Lightning-AI · Mar 21, 2024 · 84fb458 · 84fb458
1 parent 635ac44
commit 84fb458
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/thunder/tests/test_cudnn_executor.py b/thunder/tests/test_cudnn_executor.py
@@ -50,8 +50,8 @@ def grad_scaled_dot_product_attention_reference_generator(op, device, dtype, req
 
     # Non-contiguous input tensor case
     nq = make(N, n_head, E, L).permute(0, 1, 3, 2)
-    nk = make(N, n_head, E, L).permute(0, 1, 3, 2)
-    nv = make(N, n_head, E, L).permute(0, 1, 3, 2)
+    nk = make(N, n_head, E, S).permute(0, 1, 3, 2)
+    nv = make(N, n_head, E, S).permute(0, 1, 3, 2)
     yield SampleInput(nq, nk, nv, None, dropout_p=0.0, is_causal=False)
 
     # Test the scale factor which was added in torch 2.1