Update __init__.py

wcxve · Mar 26, 2024 · 4261266 · 4261266
1 parent c70a8ce
commit 4261266
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/src/elisa/__init__.py b/src/elisa/__init__.py
@@ -15,8 +15,13 @@
     DistParameter as DistParameter,
     UniformParameter as UniformParameter,
 )
-from .util import jax_enable_x64, set_cpu_cores, set_debug_nan, set_platform
+from .util import (
+    jax_debug_nans as jax_debug_nans,
+    jax_enable_x64,
+    set_cpu_cores,
+    set_jax_platform,
+)
 
 jax_enable_x64(True)
-set_platform()
+set_jax_platform('cpu')
 set_cpu_cores(4)