kamwoh commited on
Commit
226f12f
1 Parent(s): 0a52f97

fixed loading issue

Browse files
Files changed (1) hide show
  1. dreamcreature/pipeline.py +1 -1
dreamcreature/pipeline.py CHANGED
@@ -37,7 +37,7 @@ def setup_attn_processor(unet, **kwargs):
37
 
38
  def load_attn_processor(unet, filename):
39
  lora_layers = AttnProcsLayers(get_attn_processors(unet))
40
- lora_layers.load_state_dict(torch.load(filename))
41
 
42
 
43
  def convert_prompt_re(prompt: str):
 
37
 
38
  def load_attn_processor(unet, filename):
39
  lora_layers = AttnProcsLayers(get_attn_processors(unet))
40
+ lora_layers.load_state_dict(torch.load(filename, map_location='cpu'))
41
 
42
 
43
  def convert_prompt_re(prompt: str):