test (nn/sdpa): Fix when PT<2.1

Xilinx · Nov 28, 2024 · 32c94b9 · 32c94b9
1 parent 6b1e51a
commit 32c94b9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/brevitas/nn/test_sdpa.py b/tests/brevitas/nn/test_sdpa.py
@@ -37,6 +37,8 @@ def test_sdpa_fwd(self, dropout_p, is_causal, scale, enable_gqa, rand_attn_mask)
             "enable_gqa": enable_gqa,}
         if torch_version < version.parse('2.5.0'):
             del extra_kwargs["enable_gqa"]
+        if torch_version < version.parse('2.1.0'):
+            del extra_kwargs["scale"]
 
         kv_length = PAST_SEQUENCE_LENGTH + SEQUENCE_LENGTH
         m = ScaledDotProductAttention()