fix the unit tests expectations

apple · Dec 18, 2024 · 62daa78 · 62daa78
1 parent 5fdd5e3
commit 62daa78
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/coremltools/converters/mil/mil/passes/tests/test_passes.py b/coremltools/converters/mil/mil/passes/tests/test_passes.py
@@ -7471,8 +7471,8 @@ def verify_sdpa_outputs(self, example_inputs: Dict[str, torch.Tensor]):
 
         assert ops_counts[0] == 1 or ops_counts[0] == 3  # (attn_mask might be cast to bool from input fp16 dtype)
         assert ops_counts[1] == 1 or ops_counts[1] == 3  # the Q seq length is less than the default min seq length
-        assert ops_counts[2] >= 26 * 16
-        assert ops_counts[3] >= 26 * 32
+        assert ops_counts[2] >= 11 * 16  # 11 ops (without consts) per slice
+        assert ops_counts[3] >= 11 * 32
 
         predict_inputs = copy.deepcopy(example_inputs)
         if "attn_mask" in predict_inputs: