fix cpu string (#1101)

pyro-ppl · Jul 21, 2021 · b33f808 · b33f808
1 parent 339e0c0
commit b33f808
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/numpyro/util.py b/numpyro/util.py
@@ -21,7 +21,7 @@
 from jax.tree_util import tree_flatten, tree_map
 
 _DISABLE_CONTROL_FLOW_PRIM = False
-_CHAIN_RE = re.compile(r"(?<=_)\d+$")  # e.g. get '3' from 'TFRT_CPU_3'
+_CHAIN_RE = re.compile(r"\d+$")  # e.g. get '3' from 'TFRT_CPU_3'
 
 
 def set_rng_seed(rng_seed):