send weights to target device instead of CPU memory

Mapleshade20 · Aug 16, 2023 · 0815c45 · 0815c45
1 parent 8622126
commit 0815c45
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/modules/sd_models.py b/modules/sd_models.py
@@ -579,7 +579,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None):
 
     timer.record("create model")
 
-    with sd_disable_initialization.LoadStateDictOnMeta(state_dict, devices.cpu):
+    with sd_disable_initialization.LoadStateDictOnMeta(state_dict, devices.device):
         load_model_weights(sd_model, checkpoint_info, state_dict, timer)
     timer.record("load weights from state dict")