bug in more things getting batch dimension

mjo22 · Sep 16, 2024 · 0a3113a · 0a3113a
1 parent 09e84e0
commit 0a3113a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/cryojax/data/_relion/_starfile_reading.py b/src/cryojax/data/_relion/_starfile_reading.py
@@ -5,6 +5,7 @@
 from typing import Any, Callable, final, Optional
 
 import equinox as eqx
+import equinox.internal as eqxi
 import jax
 import jax.numpy as jnp
 import mrcfile
@@ -266,7 +267,11 @@ def _get_starfile_params(
             phase_shift,
         )
         ctf = (
-            eqx.filter_vmap(make_ctf, in_axes=(0, 0, 0, None, None, None, 0))(*ctf_params)
+            eqx.filter_vmap(
+                make_ctf,
+                in_axes=(0, 0, 0, None, None, None, 0),
+                out_axes=eqxi.if_mapped(0),
+            )(*ctf_params)
             if defocus_in_angstroms.ndim == 1
             else make_ctf(*ctf_params)
         )