diff --git a/sub-packages/bionemo-esm2/tests/bionemo/esm2/scripts/test_infer_esm2.py b/sub-packages/bionemo-esm2/tests/bionemo/esm2/scripts/test_infer_esm2.py index aac0ed617d..5a26e7b81e 100644 --- a/sub-packages/bionemo-esm2/tests/bionemo/esm2/scripts/test_infer_esm2.py +++ b/sub-packages/bionemo-esm2/tests/bionemo/esm2/scripts/test_infer_esm2.py @@ -32,10 +32,6 @@ from bionemo.llm.utils.callbacks import IntervalT -esm2_650m_checkpoint_path = load("esm2/650m:2.0") -esm2_3b_checkpoint_path = load("esm2/3b:2.0", source="ngc") - - # Function to check GPU memory def check_gpu_memory(threshold_gb): if torch.cuda.is_available(): @@ -140,7 +136,6 @@ def test_esm2_fine_tune_data_module_val_dataloader(data_module): @pytest.mark.parametrize("precision", ["fp32", "bf16-mixed"]) @pytest.mark.parametrize("prediction_interval", get_args(IntervalT)) -@pytest.mark.skipif(check_gpu_memory(30), reason="Skipping test due to insufficient GPU memory") def test_infer_runs( tmpdir, dummy_protein_csv, @@ -155,7 +150,7 @@ def test_infer_runs( infer_model( data_path=data_path, - checkpoint_path=esm2_650m_checkpoint_path, + checkpoint_path=load("esm2/650m:2.0"), results_path=result_dir, min_seq_length=min_seq_len, prediction_interval=prediction_interval,