Moved force cpu into cellfinder function (#430)

brainglobe · May 29, 2024 · 4ebbd46 · 4ebbd46
1 parent 4355b95
commit 4ebbd46
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/cellfinder/core/tools/system.py b/cellfinder/core/tools/system.py
@@ -1,5 +1,6 @@
 from pathlib import Path
 
+import keras
 from brainglobe_utils.general.exceptions import CommandLineInputError
 
 
@@ -80,3 +81,12 @@ def memory_in_bytes(memory_amount, unit):
         )
     else:
         return memory_amount * 10 ** supported_units[unit]
+
+
+def force_cpu():
+    """
+    Forces the CPU to be used, even if a GPU is available
+    """
+    keras.src.backend.common.global_state.set_global_attribute(
+        "torch_device", "cpu"
+    )
diff --git a/tests/core/conftest.py b/tests/core/conftest.py
@@ -1,7 +1,6 @@
 import os
 from typing import Tuple
 
-import keras.src.backend.common.global_state
 import numpy as np
 import pytest
 import torch.backends.mps
@@ -11,6 +10,7 @@
     DEFAULT_DOWNLOAD_DIRECTORY,
     download_models,
 )
+from cellfinder.core.tools.system import force_cpu
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -24,9 +24,7 @@ def set_device_arm_macos_ci():
         os.getenv("GITHUB_ACTIONS") == "true"
         and torch.backends.mps.is_available()
     ):
-        keras.src.backend.common.global_state.set_global_attribute(
-            "torch_device", "cpu"
-        )
+        force_cpu()
 
 
 @pytest.fixture(scope="session")