Fix

Author: mseeger

tests/test_model.py

@@ -1437,7 +1437,7 @@ def test_sdpa_choice(config):
     torch.set_default_dtype(torch.float16)
     config["n_layer"] = 1
     config = config_module.Config(**config)
-    enable_gqa = config["n_query_groups"] < config["n_head"]
+    enable_gqa = config.n_query_groups < config.n_head
 
     def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
         # SDPAParams gained an additional argument in PyTorch 2.5
@@ -1489,7 +1489,7 @@ def test_sdpa_choice_kv_cache(config):
     torch.set_default_dtype(torch.float16)
     config["n_layer"] = 1
     config = config_module.Config(**config)
-    enable_gqa = config["n_query_groups"] < config["n_head"]
+    enable_gqa = config.n_query_groups < config.n_head
 
     def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
         # SDPAParams gained an additional argument in PyTorch 2.5