diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 8018704f..6d6a1726 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -54,7 +54,11 @@ jobs: run: conda list - name: Build and install the package - run: pip -vv install . + run: | + if [[ "${{ runner.os }}" == "macOS" ]]; then + export OMP_PREFIX=/Users/runner/miniconda3/envs/test + fi + pip -vv install . env: CPU_ONLY: 1 diff --git a/tests/test_module.py b/tests/test_module.py index 1f553290..8a84b4e9 100644 --- a/tests/test_module.py +++ b/tests/test_module.py @@ -12,6 +12,7 @@ from torchmdnet.module import LNNP from torchmdnet.data import DataModule from torchmdnet import priors +import os from utils import load_example_args, DummyDataset @@ -30,6 +31,11 @@ def test_load_model(): @mark.parametrize("use_atomref", [True, False]) @mark.parametrize("precision", [32, 64]) def test_train(model_name, use_atomref, precision, tmpdir): + import torch + + if os.getenv("CPU_ONLY"): + torch.set_default_device("cpu") + args = load_example_args( model_name, remove_prior=not use_atomref, @@ -53,6 +59,11 @@ def test_train(model_name, use_atomref, precision, tmpdir): module = LNNP(args, prior_model=prior) - trainer = pl.Trainer(max_steps=10, default_root_dir=tmpdir, precision=args["precision"],inference_mode=False) + trainer = pl.Trainer( + max_steps=10, + default_root_dir=tmpdir, + precision=args["precision"], + inference_mode=False, + ) trainer.fit(module, datamodule) trainer.test(module, datamodule)