diff --git a/examples/whm_gr_test/swiftest_relativity.ipynb b/examples/whm_gr_test/swiftest_relativity.ipynb
index 4d44ca569..ea91ed5a2 100644
--- a/examples/whm_gr_test/swiftest_relativity.ipynb
+++ b/examples/whm_gr_test/swiftest_relativity.ipynb
@@ -40,9 +40,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%%capture\n",
-    "tstep_out = 10.0\n",
-    "sim_gr.run(tstop=1000.0, dt=0.005, tstep_out=tstep_out, integrator=\"whm\",general_relativity=True)"
+    "run_args = {\"tstop\":1000.0, \"dt\":0.005, \"tstep_out\":10.0,\"integrator\":\"whm\"}"
    ]
   },
   {
@@ -51,8 +49,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%%capture\n",
-    "sim_nogr.run(tstop=1000.0, dt=0.005, tstep_out=tstep_out, integrator=\"whm\",general_relativity=False)"
+    "sim_gr.run(**run_args,general_relativity=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sim_nogr.run(**run_args,general_relativity=False)"
    ]
   },
   {