TST: Only test on CPU (faster)

dynamicslab · Oct 16, 2024 · 3c2db9b · 3c2db9b
1 parent e4b9b95
commit 3c2db9b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/test/conftest.py b/test/conftest.py
@@ -3,6 +3,7 @@
 """
 from pathlib import Path
 
+import jax
 import numpy as np
 import pytest
 from scipy.integrate import solve_ivp
@@ -20,6 +21,8 @@
 from pysindy.utils.odes import lorenz
 from pysindy.utils.odes import lorenz_control
 
+jax.config.update("jax_platform_name", "cpu")
+
 
 def pytest_addoption(parser):
     parser.addoption(