diff --git a/tools/patch_comfy_sampler_capture.py b/tools/patch_comfy_sampler_capture.py
index e8bd8cf..e7202b7 100644
--- a/tools/patch_comfy_sampler_capture.py
+++ b/tools/patch_comfy_sampler_capture.py
@@ -13,6 +13,7 @@ new = (
     "            capture_dir = os.getenv(\"H3_SAMPLER_CAPTURE_DIR\")\n"
     "            if capture_dir:\n"
     "                os.makedirs(capture_dir, exist_ok=True)\n"
+    "                torch.cuda.synchronize(x[\"x\"].device)\n"
     "                if x[\"i\"] == 0:\n"
     "                    torch.save({\"sigmas\": sigmas.detach().cpu(), \"initial_x\": x[\"x\"].detach().cpu()}, os.path.join(capture_dir, \"initial.pt\"))\n"
     "                torch.save({\"sigma\": x[\"sigma\"].detach().cpu(), \"denoised\": x[\"denoised\"].detach().cpu(), \"x\": x[\"x\"].detach().cpu()}, os.path.join(capture_dir, f\"step_{x[\"i\"]:02d}.pt\"))\n"