From 88b5feb1edd252a6b31b48efdb16a4b680322c37 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 3 Aug 2025 22:35:24 +0300
Subject: [PATCH] Fix VRAM management (again)

---
 nodes.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/nodes.py b/nodes.py
index e232381..2a1a096 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1978,6 +1978,10 @@ class WanVideoSampler:
                     module.offload()
                 if hasattr(module, "onload"):
                     module.onload()
+            for block in transformer.blocks:
+                block.modulation = torch.nn.Parameter(block.modulation.to(device))
+            transformer.head.modulation = torch.nn.Parameter(transformer.head.modulation.to(device))
+
         elif model["manual_offloading"]:
             transformer.to(device)