MatthiasC commited on
Commit
5cb75cc
1 Parent(s): e65b549

map location of torch load to cpu

Browse files
Files changed (1) hide show
  1. server.py +1 -1
server.py CHANGED
@@ -71,7 +71,7 @@ model = Dalle.from_pretrained("minDALL-E/1.3B") # This will automatically downl
71
  # model.stage2.load_state_dict(dalle_stage_dict)
72
 
73
  # NEW METHOD
74
- model.load_state_dict(torch.load(full_dict_path))
75
  model.to(device=device)
76
 
77
  # model_clip, preprocess_clip = clip.load("ViT-B/32", device=device)
 
71
  # model.stage2.load_state_dict(dalle_stage_dict)
72
 
73
  # NEW METHOD
74
+ model.load_state_dict(torch.load(full_dict_path, map_location=torch.device('cpu')))
75
  model.to(device=device)
76
 
77
  # model_clip, preprocess_clip = clip.load("ViT-B/32", device=device)