From 1e164b6ec37eff1034c213628dfc75105922b233 Mon Sep 17 00:00:00 2001
From: u-haru <40634644+u-haru@users.noreply.github.com>
Date: Fri, 31 Mar 2023 12:52:39 +0900
Subject: [PATCH] specify device when loading state_dict

---
 library/model_util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/library/model_util.py b/library/model_util.py
index 9b4405eb..32a9c87a 100644
--- a/library/model_util.py
+++ b/library/model_util.py
@@ -866,7 +866,7 @@ def load_checkpoint_with_text_encoder_conversion(ckpt_path, device="cpu"):
 
 # TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認
 def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device='cpu', dtype=None):
-  _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path) # no need to specify device in loading state_dict
+  _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path, device)
 
   # Convert the UNet2DConditionModel model.
   unet_config = create_unet_diffusers_config(v2)