Fix

Author: mseeger

tests/test_model.py

@@ -1513,9 +1513,9 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
     try:
         with torch.device("cuda"):
             model = GPT(config)
-            model.max_seq_length = 1
+            model.max_seq_length = 16
             model.set_kv_cache(2)
-            x = torch.randint(0, 10, (2, 1), dtype=torch.int32)
+            x = torch.randint(0, 10, (2, 10), dtype=torch.int32)
     except torch.cuda.OutOfMemoryError:
         # best effort, if the GPU can load it
         pytest.xfail()