Pass the correct quantities to run_vmap (#62)

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
GalacticDynamics · Jan 18, 2024 · fb2a422 · fb2a422
1 parent 0536100
commit fb2a422
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/galax/dynamics/mockstream/_mockstream_generator.py b/src/galax/dynamics/mockstream/_mockstream_generator.py
@@ -141,7 +141,7 @@ def run(
         )
 
         if vmapped:
-            lead_arm_qp, trail_arm_qp = self._run_vmap(ts, prog_w0, prog_mass)
+            lead_arm_qp, trail_arm_qp = self._run_vmap(ts, mock0_lead, mock0_trail)
         else:
             lead_arm_qp, trail_arm_qp = self._run_scan(ts, mock0_lead, mock0_trail)