Upload 2 files

efficient_inference_for_square_image.py

@@ -335,7 +335,7 @@ def main_process(opt, composite_image=None, mask=None):
 
     model = build_model(opt).to(opt.device)
 
-    load_dict = torch.load(opt.pretrained)['model']
+    load_dict = torch.load(opt.pretrained, map_location='cpu')['model']
     for k in load_dict.keys():
         if k not in model.state_dict().keys():
             print(f"Skip {k}")

inference_for_arbitrary_resolution_image.py

@@ -327,7 +327,7 @@ def main_process(opt, composite_image=None, mask=None):
 
     model = build_model(opt).to(opt.device)
 
-    load_dict = torch.load(opt.pretrained)['model']
+    load_dict = torch.load(opt.pretrained, map_location='cpu')['model']
     for k in load_dict.keys():
         if k not in model.state_dict().keys():
             print(f"Skip {k}")