Update nodes_model_loading.py

2025-08-29 21:21:24 +03:00
parent 3a290cfd15
commit 7dd0e1ac61
1 changed files with 1 additions and 0 deletions
@@ -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