fixed a > and >= in masking decision / added to test suite

tjgalvin · Jun 21, 2024 · 5899bc3 · 5899bc3
1 parent 263b261
commit 5899bc3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/flint/masking.py b/flint/masking.py
@@ -102,7 +102,7 @@ def consider_beam_mask_round(
 
     return mask_rounds is not None and (
         (isinstance(mask_rounds, str) and mask_rounds.lower() == "all")
-        or (isinstance(mask_rounds, int) and current_round > mask_rounds)
+        or (isinstance(mask_rounds, int) and current_round >= mask_rounds)
         or (isinstance(mask_rounds, Iterable) and current_round in mask_rounds)  # type: ignore
     )
 

diff --git a/tests/test_masking.py b/tests/test_masking.py
@@ -39,6 +39,7 @@ def test_consider_beam_masking_round():
     assert consider_beam_mask_round(
         current_round=3, mask_rounds=1, allow_beam_masks=True
     )
+    assert consider_beam_mask_round(current_round=1, mask_rounds=1)
 
 
 def test_minimum_boxcar_artefact():