From 7dd0e1ac617534d70cd1c11483e62716143856b1 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 29 Aug 2025 21:21:24 +0300
Subject: [PATCH] Update nodes_model_loading.py

---
 nodes_model_loading.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nodes_model_loading.py b/nodes_model_loading.py
index 52b12f0..a0c747b 100644
--- a/nodes_model_loading.py
+++ b/nodes_model_loading.py
@@ -824,6 +824,7 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None,
 
         load_device = transformer_load_device
         if block_swap_args is not None:
+            load_device = device
             if block_idx is not None:
                 if block_idx >= len(transformer.blocks) - block_swap_args.get("blocks_to_swap", 0):
                     load_device = offload_device