send weights to target device instead of CPU memory

This commit is contained in:
AUTOMATIC1111 2023-08-16 10:44:17 +03:00
parent 86221269f9
commit 0815c45bcd

View File

@ -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")