From 3c2db9b5a8e3613a00ad5c54ecec108f45ae63ce Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Wed, 16 Oct 2024 15:31:17 -0700
Subject: [PATCH] TST: Only test on CPU (faster)

---
 test/conftest.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/conftest.py b/test/conftest.py
index 59fb61f74..ead11adf1 100644
--- 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(