From 9b1b0505b308562475defaf80fb71c86eeb7516b Mon Sep 17 00:00:00 2001
From: Igor Tatarnikov <igor.tatarnikov@gmail.com>
Date: Thu, 30 May 2024 16:54:55 +0100
Subject: [PATCH] Try using the CI variable instead

---
 tests/conftest.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 10bfe0f..da03391 100644
--- 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()