Remove vmap from kabsch

plainerman · plainerman · commit 2ef724bdc5b1 · 2024-05-07T14:53:04.000+02:00
diff --git a/utils/rmsd.py b/utils/rmsd.py
@@ -43,7 +43,6 @@ def kabsch_align(P, Q):
 
 
 @jax.jit
-@jax.vmap
 def kabsch_rmsd(P, Q):
     P_aligned, Q_aligned = kabsch_align(P, Q)
     return jnp.sqrt(jnp.sum(jnp.square(P_aligned - Q_aligned)) / P.shape[0])
