From 663534f5dc2f177be8df23fe0c425b93b0c5f61a Mon Sep 17 00:00:00 2001
From: Daniel Maddern <daniel.maddern@illuvium.io>
Date: Thu, 13 Aug 2026 13:13:29 +0700
Subject: [PATCH] Synchronize H3 sampler captures

---
 tools/patch_comfy_sampler_capture.py | 1 +
 1 file changed, 1 insertion(+)

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"