fixing test issues on OSX

torchmd · Jan 22, 2025 · b9e24c8 · b9e24c8
1 parent aacd6c7
commit b9e24c8
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git 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
@@ -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)