From dcc75ecfe5e9d5466284d7cf08d86cd4c24a6544 Mon Sep 17 00:00:00 2001 From: John St John Date: Tue, 12 Nov 2024 17:34:08 +0000 Subject: [PATCH] reduce atol check for geneformer --- .../tests/bionemo/geneformer/test_stop_and_go.py | 1 + .../src/bionemo/testing/harnesses/stop_and_go.py | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_stop_and_go.py b/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_stop_and_go.py index 3820844131..a33a6ff1ad 100644 --- a/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_stop_and_go.py +++ b/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_stop_and_go.py @@ -97,6 +97,7 @@ class TestGeneformerStopAndGo(stop_and_go.StopAndGoHarness): limit_val_batches: int = 2 lr: float = 1e-4 precision: Literal["16-mixed", "bf16-mixed", "32"] = MODEL_PRECISION + train_output_atol: float = 2e-2 @override @classmethod diff --git a/sub-packages/bionemo-testing/src/bionemo/testing/harnesses/stop_and_go.py b/sub-packages/bionemo-testing/src/bionemo/testing/harnesses/stop_and_go.py index d811c997c4..afc12a648a 100644 --- a/sub-packages/bionemo-testing/src/bionemo/testing/harnesses/stop_and_go.py +++ b/sub-packages/bionemo-testing/src/bionemo/testing/harnesses/stop_and_go.py @@ -106,6 +106,8 @@ class StopAndGoHarness(ABC): limit_val_batches: int lr: float = 1e-4 precision: Literal["16-mixed", "bf16-mixed", "32"] + train_output_atol: float = 1e-3 + other_output_atol: float = 1e-4 # class variables that will be setup in setUpClass tempdir: tempfile.TemporaryDirectory @@ -336,9 +338,9 @@ def test_stop_and_go_consistency(self, callback_type): assert interrupted_callback.data, f"No data found for {callback_type}" if callback_type == testing_callbacks.TrainOutputCallback: - atol = 1e-3 + atol = self.train_output_atol else: - atol = 1e-4 + atol = self.other_output_atol recursive_assert_approx_equal(interrupted_callback.data, continuous_callback.data, atol=atol)