psd_inv_cholesky -> psd_inv

google-deepmind · Nov 6, 2023 · 19fe893 · 19fe893
1 parent c2f345d
commit 19fe893
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ferminet/curvature_tags_and_blocks.py b/ferminet/curvature_tags_and_blocks.py
@@ -24,7 +24,7 @@
 Scalar = kfac_jax.utils.Scalar
 Numeric = kfac_jax.utils.Numeric
 
-vmap_psd_inv_cholesky = jax.vmap(kfac_jax.utils.psd_inv_cholesky, (0, None), 0)
+vmap_psd_inv = jax.vmap(kfac_jax.utils.psd_inv, (0, None), 0)
 vmap_matmul = jax.vmap(jnp.matmul, in_axes=(0, 0), out_axes=0)