Changed method name

Routhleck · Oct 12, 2023 · 405ec4f · 405ec4f
1 parent fa5fbe3
commit 405ec4f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/catalax/__init__.py b/catalax/__init__.py
@@ -19,7 +19,7 @@ def set_host_count(n: int = 1):
     """
     import numpyro
 
-    numpyro.set_host_count(n)
+    numpyro.set_host_device_count(n)
 
 
 def set_platform(platform: str = "cpu"):