diff --git a/pyproject.toml b/pyproject.toml
index 068ef47..e91ef78 100644
--- 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
index 0c1bcf9..b61cb63 100644
--- 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()