Use force_cpu function from cellfinder

brainglobe · May 31, 2024 · 0ab452a · 0ab452a
1 parent 066fbd9
commit 0ab452a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 10 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -147,5 +147,7 @@ commands =
 description =
     Run tests
     coredev: Run tests with the development version of cellfinder
-passenv = GITHUB_*
+passenv =
+    CI
+    GITHUB_*
 """
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -2,8 +2,7 @@
 
 import pytest
 import torch
-
-# from cellfinder.core.tools.system import force_cpu
+from cellfinder.core.tools.system import force_cpu
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -22,10 +21,4 @@ def set_device_arm_macos_ci():
         os.getenv("GITHUB_ACTIONS") == "true"
         and torch.backends.mps.is_available()
     ):
-        import keras
-
-        keras.src.backend.common.global_state.set_global_attribute(
-            "torch_device", "cpu"
-        )
-
-        # force_cpu()
+        force_cpu()