diff --git a/tests/test_against_fpm.py b/tests/test_against_fpm.py
index 6d17939..e02eff4 100644
--- a/tests/test_against_fpm.py
+++ b/tests/test_against_fpm.py
@@ -76,7 +76,7 @@ def test_nbody_absolute(simulation_config, initial_conditions,
                    a=lpt_scale_factor,
                    order=order)
 
-    ode_fn = ODETerm(make_diffrax_ode(cosmo, mesh_shape))
+    ode_fn = ODETerm(make_diffrax_ode(mesh_shape))
 
     solver = Dopri5()
     controller = PIDController(rtol=1e-8,
@@ -122,7 +122,7 @@ def test_nbody_relative(simulation_config, initial_conditions,
     dx, p, _ = lpt(cosmo, initial_conditions, a=lpt_scale_factor, order=order)
 
     ode_fn = ODETerm(
-        make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False))
+        make_diffrax_ode(mesh_shape, paint_absolute_pos=False))
 
     solver = Dopri5()
     controller = PIDController(rtol=1e-9,