Try using the CI variable instead

brainglobe · May 30, 2024 · 9b1b050 · 9b1b050
1 parent 9a70828
commit 9b1b050
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,5 +1,4 @@
 import os
-import sys
 
 import pytest
 import torch
@@ -19,10 +18,11 @@ def set_device_arm_macos_ci():
     GitHub runners. This is to avoid the following error:
     https://discuss.pytorch.org/t/mps-back-end-out-of-memory-on-github-action/189773/5
     """
-    if sys.platform == "darwin" and torch.backends.mps.is_available():
+    if os.getenv("CI") == "true" and torch.backends.mps.is_available():
         import keras
 
         keras.src.backend.common.global_state.set_global_attribute(
             "torch_device", "cpu"
         )
+
         # force_cpu()