cuda device

app.py

@@ -66,7 +66,8 @@ def save_story_results(images, video_len=4, n_candidates=1, mask=None):
 
 
 def main(args):
-    device = 'cuda:0'
+    #device = 'cuda:0'
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
     model_url = 'https://drive.google.com/u/1/uc?id=1KAXVtE8lEE2Yc83VY7w6ycOOMkdWbmJo&export=sharing'
 

dalle/models/__init__.py

@@ -1193,9 +1193,9 @@ class StoryDalle(Dalle):
             print("Loading model from pretrained checkpoint %s" % args.model_name_or_path)
             # model.from_ckpt(args.model_name_or_path)
             try:
-                model.load_state_dict(torch.load(args.model_name_or_path)['state_dict'])
+                model.load_state_dict(torch.load(args.model_name_or_path, map_location=torch.device('cpu'))['state_dict'])
             except KeyError:
-                model.load_state_dict(torch.load(args.model_name_or_path)['model_state_dict'])
+                model.load_state_dict(torch.load(args.model_name_or_path, map_location=torch.device('cpu'))['model_state_dict'])
         else:
             model = cls(config_update)
             print(model.cross_attention_idxs)